React Compiler, How Does It Work? [1] In the previous article, we looked at the overall structure of React Compiler.
Before examining how the compiler works, let’s first take a look at useMemoCache, which was mentioned frequently.
useMemoCache implementation by josephsavona · Pull Request #25143 · facebook/react · GitHub
Summary context: I'm still a noob to this codebase so this started as a hacky implementation to get feedback. Initial implementation of useMemoCache that aims to at least get the basics correct, ev...
https://github.com/facebook/react/pull/25143This is the initial implementation PR for useMemoCache. If you’re curious, take a look.
Let’s first skim through the entire code and then examine it part by part
// react-reconciler/src/ReactFiberHooks.js
function useMemoCache(size: number): Array<any> {
let memoCache = null;
// Fast-path, load memo cache from wip fiber if already prepared
let updateQueue: FunctionComponentUpdateQueue | null =
(currentlyRenderingFiber.updateQueue: any);
if (updateQueue !== null) {
memoCache = updateQueue.memoCache;
}
// Otherwise clone from the current fiber
if (memoCache == null) {
const current: Fiber | null = currentlyRenderingFiber.alternate;
if (current !== null) {
const currentUpdateQueue: FunctionComponentUpdateQueue | null =
(current.updateQueue: any);
if (currentUpdateQueue !== null) {
const currentMemoCache: ?MemoCache = currentUpdateQueue.memoCache;
if (currentMemoCache != null) {
memoCache = {
// When enableNoCloningMemoCache is enabled, instead of treating the
// cache as copy-on-write, like we do with fibers, we share the same
// cache instance across all render attempts, even if the component
// is interrupted before it commits.
//
// If an update is interrupted, either because it suspended or
// because of another update, we can reuse the memoized computations
// from the previous attempt. We can do this because the React
// Compiler performs atomic writes to the memo cache, i.e. it will
// not record the inputs to a memoization without also recording its
// output.
//
// This gives us a form of "resuming" within components and hooks.
//
// This only works when updating a component that already mounted.
// It has no impact during initial render, because the memo cache is
// stored on the fiber, and since we have not implemented resuming
// for fibers, it's always a fresh memo cache, anyway.
//
// However, this alone is pretty useful — it happens whenever you
// update the UI with fresh data after a mutation/action, which is
// extremely common in a Suspense-driven (e.g. RSC or Relay) app.
data: enableNoCloningMemoCache
? currentMemoCache.data
: // Clone the memo cache before each render (copy-on-write)
currentMemoCache.data.map(array => array.slice()),
index: 0,
};
}
}
}
}
// Finally fall back to allocating a fresh instance of the cache
if (memoCache == null) {
memoCache = {
data: [],
index: 0,
};
}
if (updateQueue === null) {
updateQueue = createFunctionComponentUpdateQueue();
currentlyRenderingFiber.updateQueue = updateQueue;
}
updateQueue.memoCache = memoCache;
let data = memoCache.data[memoCache.index];
if (data === undefined) {
data = memoCache.data[memoCache.index] = new Array(size);
for (let i = 0; i < size; i++) {
data[i] = REACT_MEMO_CACHE_SENTINEL;
}
} else if (data.length !== size) {
// TODO: consider warning or throwing here
if (__DEV__) {
console.error(
'Expected a constant size argument for each invocation of useMemoCache. ' +
'The previous cache was allocated with size %s but size %s was requested.',
data.length,
size,
);
}
}
memoCache.index++;
return data;
}
Now let’s examine the hook in order.
Fast path
The first operation is to find memoCache from the updateQueue of the fiber currently being rendered.
If it has been called (rendered) before, it will exist. This can prevent unnecessary cache allocation.
let memoCache = null;
// Fast-path, load memo cache from wip fiber if already prepared
let updateQueue: FunctionComponentUpdateQueue | null =
(currentlyRenderingFiber.updateQueue: any);
if (updateQueue !== null) {
memoCache = updateQueue.memoCache;
}
If there’s no cache? Clone cache from alternate fiber
If there’s no memoCache in the updateQueue of the fiber currently being rendered, check the alternate fiber.
In React’s Fiber architecture, each Fiber node belongs to either the ‘current’ tree or the ‘workInProgress’ tree.
The ‘current’ tree represents the state of components currently rendered on the screen, and the ‘workInProgress’ tree represents the state that React is trying to apply updates to.
The alternate fiber (Alternate Fiber) mentioned here is the alternate of currentlyRenderingFiber, i.e., the workInProgress tree, so it points to the Fiber of the current tree.\
Clone the cache from this alternate fiber (current) for use.
I’ve temporarily removed the conditional statements to see the flow of logic. You can check the original text above.
let memoCache = null;
// ...
// if cache is not exist, clone cache from alternate fiber
if (memoCache == null) {
const current = currentlyRenderingFiber.alternate; // Get the alternate fiber (current).
const currentUpdateQueue = current.updateQueue; // Get the updateQueue of the alternate fiber.
const currentMemoCache: ?MemoCache = currentUpdateQueue.memoCache; //Get the memoCache of the alternate fiber.
memoCache = {
data: enableNoCloningMemoCache
? currentMemoCache.data
: // Clone the memo cache before each render (copy-on-write)
currentMemoCache.data.map(array => array.slice()),
index: 0,
};
}
At this point, the behavior differs depending on the enableNoCloningMemoCache option.
The enableNoCloningMemoCache flag determines whether to use the cache data as is or make a shallow copy.
- If enabled, use the cache data as is. This can reduce memory usage but carries the risk of cache data being changed.
- If disabled, use a shallow copy of the cache data. This allows safe use of cache data from previous rendering, but may increase memory usage.
[Experiment] Reuse memo cache after interruption by acdlite · Pull Request #28878 · facebook/react · GitHub
Adds an experimental feature flag to the implementation of useMemoCache, the internal cache used by the React Compiler (Forget). When enabled, instead of treating the cache as copy-on-write, like w...
https://github.com/facebook/react/pull/28878There’s a long comment attached to this part, let’s take a look at it too.
When enableNoCloningMemoCache is enabled, instead of treating the
cache as copy-on-write, like we do with fibers, we share the same
cache instance across all render attempts, even if the component
is interrupted before it commits.
If an update is interrupted, either because it suspended or
because of another update, we can reuse the memoized computations
from the previous attempt. We can do this because the React
Compiler performs atomic writes to the memo cache, i.e. it will
not record the inputs to a memoization without also recording its
output.
This gives us a form of "resuming" within components and hooks.
This only works when updating a component that already mounted.
It has no impact during initial render, because the memo cache is
stored on the fiber, and since we have not implemented resuming
for fibers, it's always a fresh memo cache, anyway.
However, this alone is pretty useful — it happens whenever you
update the UI with fresh data after a mutation/action, which is
extremely common in a Suspense-driven (e.g. RSC or Relay) app.
What does it mean that This is possible because the React Compiler performs atomic writes to the memo cache. In other words, it doesn't record inputs for memoization without recording outputs?
🤷♂️ Why does atom come up here?
First, the word ‘atomic’ is used to mean ‘indivisible’. (Let’s not say atoms can be divided too)
In computer science, an ‘atomic operation’ means an ‘indivisible operation’, that is, an operation that is either completely executed or not executed at all.
Let’s take an example.
Think about a bank transfer. When transferring 100,000 won from account A to account B, this operation is done in two steps.
- The operation of subtracting 100,000 won from account A
- The operation of adding 100,000 won to account B
What happens if there’s a problem between these processes and the process is interrupted?
100,000 won disappears from account A, but account B doesn’t receive it. To prevent this situation, we need ‘atomic operations’.
In other words, the transfer operation should either be executed completely at once or not executed at all.
The memoization process of React Compiler uses a similar concept. Let’s explain with an example.
Let’s compile the code below.
function Component({ active }) {
let color: string;
if (active) {
color = "red";
} else {
color = "blue";
}
return <div styles={{ color }}>hello world</div>;
}
When compiled, it will be transformed as follows.
function Component(t0) {
const $ = _c(2);
const { active } = t0;
let color;
if (active) {
color = "red";
} else {
color = "blue";
}
let t1;
if ($[0] !== color) {
t1 = (
<div styles={{ color }}>hello world</div>
);
$[0] = color;
$[1] = t1;
} else {
t1 = $[1];
}
return t1;
}
Now let’s look at the part we should focus on.
if ($[0] !== color) {
t1 = (
<div styles={{ color }}>hello world</div>
);
$[0] = color;
$[1] = t1;
} else {
t1 = $[1];
}
Let’s briefly explain how it works.
$ means memoCache. The compiler stores the previously stored color value in $[0], and the previously rendered result in $[1].
if ($[0] !== color) compares the cached color value with the current color value. If they’re different, it means we need to create a new element for the new color.
At this time, it creates a new element, stores the new color value in $[0], and stores the new element in $[1].
If not, it uses the previously rendered result.
The important point here is that $[0] = color; and $[1] = t1; happen consecutively in the same block.
This is exactly ‘atomic writing’.
These two lines are treated as one atomic operation, so when the color value is written to the cache, the corresponding element is also definitely written to the cache. There is no intermediate state between them.
Now we know what React Compiler’s atomic writing means.
Let’s go back to the original content.
If enableNoCloningMemoCache is false, it uses the cache from the previous rendering by copying and modifying this copy (copy-on-write) for each rendering attempt.
However, this method can increase memory usage because it copies the cache every time rendering is interrupted and resumed.
If enableNoCloningMemoCache is true, it shares the same cache instance across all rendering attempts.
In other words, even if rendering is interrupted (suspended/interrupted) and resumed, it can use the cache from the previous rendering attempt as is.
This is a big advantage in terms of memory management.
For this behavior to be possible, the React compiler must perform atomic writes to the memo cache.
For example, if it wasn’t guaranteed, the following problem would have occurred:
- Rendering attempt ‘A’ stores the color value ‘red’ in the cache.
- However, rendering is interrupted before recording the element in the cache.
- When rendering attempt ‘B’ starts, at this point the cache has input for color, but no output.
This will cause a cache inconsistency problem and this will lead to inconsistency in rendering results.
Inconsistency in rendering results due to state… Doesn’t this sound familiar? It’s similar to the tearing problem that comes up when explaining ‘useSyncExternalStore’.
What is tearing? · reactwg/react-18 · Discussion #69 · GitHub
What is tearing?
https://github.com/reactwg/react-18/discussions/69However, if ‘atomic writing’ is guaranteed, having a cache for color ensures that there must also be a cache for the element.
Whew, we’ve come quite far. Now let’s go back to useMemoCache.
We came this far while looking at the branch according to the enableNoCloningMemoCache condition.
So is enableNoCloningMemoCache true or false!
// react/shared/ReactFeatureFlags.js
// Test this at Meta before enabling.
export const enableNoCloningMemoCache = false;
At the current point, enableNoCloningMemoCache is set to false.
Let’s go back to useMemoCache.
If there’s no cache? Create a new cache
Now finally, if there’s no cache, create a new cache.
if (memoCache == null) {
memoCache = {
data: [],
index: 0,
};
}
Cache allocation
If there’s no updateQueue, create a new one and allocate memoCache.
if (updateQueue === null) {
updateQueue = createFunctionComponentUpdateQueue();
currentlyRenderingFiber.updateQueue = updateQueue;
}
updateQueue.memoCache = memoCache;
createFunctionComponentUpdateQueue returns a basic queue object.
It’s divided according to the enableUseMemoCacheHook flag, which is currently set to true.
// NOTE: defining two versions of this function to avoid size impact when this feature is disabled.
// Previously this function was inlined, the additional `memoCache` property makes it not inlined.
let createFunctionComponentUpdateQueue: () => FunctionComponentUpdateQueue;
if (enableUseMemoCacheHook) {
createFunctionComponentUpdateQueue = () => {
return {
lastEffect: null,
events: null,
stores: null,
memoCache: null,
};
};
} else {
createFunctionComponentUpdateQueue = () => {
return {
lastEffect: null,
events: null,
stores: null,
};
};
}
Return cache data
Finally, return the cache data.
Get data from the cache,
let data = memoCache.data[memoCache.index];
If there’s no data, allocate new data and initialize it with REACT_MEMO_CACHE_SENTINEL.
REACT_MEMO_CACHE_SENTINEL represents the state where memoization has not been done.
Because cache is used in Array form for each call, it’s initialized as an Array.
if (data === undefined) {
data = memoCache.data[memoCache.index] = new Array(size);
for (let i = 0; i < size; i++) {
data[i] = REACT_MEMO_CACHE_SENTINEL;
}
}
If data is not undefined, that is, if cache data exists, check if the length of data matches the requested size.
The reason for checking is that if the length of the cache data used in the previous rendering is different from the length of the cache data to be used in the current rendering, problems may occur.
else if (data.length !== size) {
if (__DEV__) {
console.error(
'Expected a constant size argument for each invocation of useMemoCache. ' +
'The previous cache was allocated with size %s but size %s was requested.',
data.length,
size,
);
}
}
Finally, increase the cache index and return the data.
memoCache.index++;
return data;
It’s similar to the implementation in ‘react-compiler-runtime’ that we looked at in the previous edition, but there are some parts that have been modified to fit Fiber better.
Difference between useMemo and useMemoCache
useMemo, which is familiar to us and will now be forgotten. Due to the similar names, you might wonder about the difference.
Let’s briefly review the implementation of useMemo.
function mountMemo<T>(
nextCreate: () => T,
deps: Array<mixed> | void | null,
): T {
const hook = mountWorkInProgressHook();
const nextDeps = deps === undefined ? null : deps;
const nextValue = nextCreate();
hook.memoizedState = [nextValue, nextDeps];
return nextValue;
}
function updateMemo<T>(
nextCreate: () => T,
deps: Array<mixed> | void | null,
): T {
const hook = updateWorkInProgressHook();
const nextDeps = deps === undefined ? null : deps; //
const prevState = hook.memoizedState;
if (prevState !== null) {
// Assume these are defined. If they're not, areHookInputsEqual will warn.
if (nextDeps !== null) {
const prevDeps: Array<mixed> | null = prevState[1];
if (areHookInputsEqual(nextDeps, prevDeps)) { // If the previous deps and the current deps are the same, return the previous value
return prevState[0];
}
}
}
const nextValue = nextCreate(); // create a new value
hook.memoizedState = [nextValue, nextDeps]; // update the memoizedState
return nextValue; // return the new value
}
Both are hooks for memoization, but they’re a bit different. useMemo is a hook that developers use directly, receiving a dependency array and creating a new value only when the dependency changes.
useMemo is responsible for memoization through the dependency array.
On the other hand, useMemoCache is a hook used internally by the compiler, doesn’t receive a dependency array, and its main purpose is to manage the cache.
Therefore, the responsibility for memoization lies with the compiler, not useMemoCache.
Test case for useMemoCache
Now let’s look at example code using useMemoCache.
There are many good examples in the test code, so looking at the test code will help understand how it works.
// Some expensive processing...
function someExpensiveProcessing(t) {
Scheduler.log(`Some expensive processing... [${t}]`);
return t;
}
// function to trigger suspense
function useWithLog(t, msg) {
try {
return React.use(t);
} catch (x) {
Scheduler.log(`Suspend! [${msg}]`);
throw x;
}
}
function Data({chunkA, chunkB}) {
const a = someExpensiveProcessing(useWithLog(chunkA, 'chunkA'));
const b = useWithLog(chunkB, 'chunkB');
return (
<>
{a}
{b}
</>
);
}
function Input() {
const [input, _setText] = useState('');
return input;
}
function App({chunkA, chunkB}) {
return (
<>
<div>
Input: <Input />
</div>
<div>
Data: <Data chunkA={chunkA} chunkB={chunkB} />
</div>
</>
);
}
These are the components in this form, let’s test with the compilation results.
The Data component receives chunkA and chunkB, performs expensive calculations on chunkA, and returns chunkB as is.
// react-reconciler/src/__tests__/useMemoCache-test.js
test('reuses computations from suspended/interrupted render attempts during an update', async () => {
// This test demonstrates the benefit of a shared memo cache. By "shared" I
// mean multiple concurrent render attempts of the same component/hook use
// the same cache. (When the feature flag is off, we don't do this — the
// cache is copy-on-write.)
function Data(t0) {
const $ = useMemoCache(5);
const {chunkA, chunkB} = t0;
const t1 = useWithLog(chunkA, 'chunkA');
let t2;
if ($[0] !== t1) {
t2 = someExpensiveProcessing(t1);
$[0] = t1;
$[1] = t2;
} else {
t2 = $[1];
}
const a = t2;
const b = useWithLog(chunkB, 'chunkB');
let t3;
if ($[2] !== a || $[3] !== b) {
t3 = (
<>
{a}
{b}
</>
);
$[2] = a;
$[3] = b;
$[4] = t3;
} else {
t3 = $[4];
}
return t3;
}
let setInput;
function Input() {
const [input, _set] = useState('');
setInput = _set;
return input;
}
function App(t0) {
const $ = useMemoCache(4);
const {chunkA, chunkB} = t0;
let t1;
if ($[0] === Symbol.for('react.memo_cache_sentinel')) {
t1 = (
<div>
Input: <Input />
</div>
);
$[0] = t1;
} else {
t1 = $[0];
}
let t2;
if ($[1] !== chunkA || $[2] !== chunkB) {
t2 = (
<>
{t1}
<div>
Data: <Data chunkA={chunkA} chunkB={chunkB} />
</div>
</>
);
$[1] = chunkA;
$[2] = chunkB;
$[3] = t2;
} else {
t2 = $[3];
}
return t2;
}
// create Resolved Promise
function createInstrumentedResolvedPromise(value) {
return {
then() {},
status: 'fulfilled',
value,
};
}
// create Pending Promise
function createDeferred() {
let resolve;
const p = new Promise(res => {
resolve = res;
});
p.resolve = resolve;
return p;
}
Let’s run the test based on the compiled code.
First rendering
In the initial rendering, it receives chunkA and chunkB and renders the Data component.
createInstrumentedResolvedPromise creates a Resolved promise. This means data that has already been received.
At this time, a log for the expensive process hanging on A1 is printed, and the Data component returns A1B1.
// Initial render. We pass the data in as two separate "chunks" to simulate a stream (e.g. RSC).
const root = ReactNoop.createRoot();
const initialChunkA = createInstrumentedResolvedPromise('A1'); // create a promise resolved to A1
const initialChunkB = createInstrumentedResolvedPromise('B1'); // create a promise resolved to B1
await act(() =>
root.render(<App chunkA={initialChunkA} chunkB={initialChunkB} />), // initial render
);
assertLog(['Some expensive processing... [A1]']); //
expect(root).toMatchRenderedOutput(
<>
<div>Input: </div>
<div>Data: A1B1</div>
</>,
);
UI update during transition
const updatedChunkA = createDeferred();
const updatedChunkB = createDeferred();
The createDeferred function is used to create two Promise objects called updatedChunkA and updatedChunkB. These objects are used to load data later.
await act(() => {
React.startTransition(() => {
root.render(<App chunkA={updatedChunkA} chunkB={updatedChunkB} />);
});
});
Use React.startTransition to start a UI update. Transition is scheduled with low priority and executed after other tasks are completed.
At this time, render the App component by passing updatedChunkA and updatedChunkB.
assertLog(['Suspend! [chunkA]']);
const t1 = useWithLog(chunkA, 'chunkA');
useWithLog function executes useWithLog for chunkA, and because it’s pending, Suspense is triggered by use and ‘Suspend! [chunkA]’ is printed.
await act(() => updatedChunkA.resolve('A2'));
Resolve updatedChunkA to ‘A2’.
const t1 = useWithLog(chunkA, 'chunkA');
let t2;
if ($[0] !== t1) {
t2 = someExpensiveProcessing(t1);
$[0] = t1;
$[1] = t2;
} else {
t2 = $[1];
}
const a = t2;
const b = useWithLog(chunkB, 'chunkB');
Since the data is ready, rendering resumes and useWithLog is executed.
At this time, t1 becomes ‘A2’, and because it’s different from ‘A1’ stored in $[0] previously, the expensive process is executed.
Accordingly, ‘Some expensive processing… [A2]’ is printed.
Immediately after, useWithLog is executed, and since chunkB is not yet resolved, ‘Suspend! [chunkB]’ is printed.
Because rendering is interrupted again, it shows the initial UI as is.
expect(root).toMatchRenderedOutput(
<>
<div>Input: </div>
<div>Data: A1B1</div>
</>,
);
Update other parts during update transition
// While waiting for the data to finish loading, update a different part of the screen. This interrupts the refresh transition.
//
// In a real app, this might be an input or hover event.
await act(() => setInput('hi!')); // update the input value to 'hi!'
Change the input value to ‘hi!’ using setInput. This update interrupts the transition and starts a new update.
// Once the input has updated, we go back to rendering the transition.
if (gate(flags => flags.enableNoCloningMemoCache)) {
// We did not have process the first chunk again. We reused the computation from the earlier attempt.
assertLog(['Suspend! [chunkB]']);
} else {
// Because we clone/reset the memo cache after every aborted attempt, we must process the first chunk again.
assertLog(['Some expensive processing... [A2]', 'Suspend! [chunkB]']);
}
If enableNoCloningMemoCache is activated, it reuses the previously calculated data, and useWithLog for chunkB is executed, printing ‘Suspend! [chunkB]‘.
However, if it’s deactivated, because cache is managed through shallow copying, in the case of memoCache managed as a 2D array, $[0] !== t1 becomes true and the expensive process is executed.
expect(root).toMatchRenderedOutput(
<>
<div>Input: hi!</div>
<div>Data: A1B1</div>
</>,
);
Since chunkB is still not resolved, it shows the previous rendering result.
await act(() => updatedChunkB.resolve('B2')); // resolve chunkB to B2
Resolve updatedChunkB to ‘B2’. Rendering resumes again.
if (gate(flags => flags.enableNoCloningMemoCache)) {
// We did not have process the first chunk again. We reused the computation from the earlier attempt.
assertLog([]);
} else {
// Because we clone/reset the memo cache after every aborted attempt, we must process the first chunk again.
//
// That's three total times we've processed the first chunk, compared to just once when enableNoCloningMemoCache is on.
assertLog(['Some expensive processing... [A2]']);
}
expect(root).toMatchRenderedOutput( // check the rendering result
<>
<div>Input: hi!</div>
<div>Data: A2B2</div>
</>,
);
If the flag is activated, rendering completes without logs as it reuses previously calculated data.
However, if it’s deactivated, ‘Some expensive processing… [A2]’ is printed as it doesn’t reuse previously calculated data.
As a result, the Data component returns A2B2.
If deactivated, a total of 3 expensive calculations occur.
But if activated, it ends with 1 calculation. That’s quite a big difference.
Currently it’s set to false, but if it changes to true in the future, there will be another improvement in performance.
So now we’ve confirmed the operation of useMemoCache even through test code.
Conclusion
In this article, we looked at useMemoCache.
useMemoCache is a hook used internally by the compiler, and its main purpose is to manage cache for memoization.
Although it’s still an experimental feature, it seems to be considering the feature of sharing cache as well.
Our view of the compiler has broadened more than yesterday.
Now that we know how memoization is implemented basically, it would be good to look at how the compiler actually works next time.
I always worry about how to end the article.
This time, I’ll just end it like this.
Goodbye!