How does useImperativeHandle() work internally?

Have you used useImperativeHandle() before ? Let’s figure out how it works intenrally.

Usage

Here is the official example usage.

jsx
function FancyInput(props, ref) {
const inputRef = useRef();
useImperativeHandle(ref, () => ({
focus: () => {
inputRef.current.focus();
},
}));
return <input ref={inputRef} />;
}
FancyInput = forwardRef(FancyInput);
jsx
function FancyInput(props, ref) {
const inputRef = useRef();
useImperativeHandle(ref, () => ({
focus: () => {
inputRef.current.focus();
},
}));
return <input ref={inputRef} />;
}
FancyInput = forwardRef(FancyInput);

By above code, we can attatch a ref to FancyInput now.

jsx
function App() {
const ref = useRef();
const focus = useCallback(() => {
ref.current?.focus();
}, []);
return (
<div>
<FancyInput ref={inputRef} />
<button onClick={focus} />
</div>
);
}
jsx
function App() {
const ref = useRef();
const focus = useCallback(() => {
ref.current?.focus();
}, []);
return (
<div>
<FancyInput ref={inputRef} />
<button onClick={focus} />
</div>
);
}

Looks straightforward, but why we do this?

what if we just update ref.current?

Rather than useImperativeHandle(), what if we just update ref.current? like below.

js
function FancyInput(props, ref) {
const inputRef = useRef();
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
return <input ref={inputRef} />;
}
js
function FancyInput(props, ref) {
const inputRef = useRef();
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
return <input ref={inputRef} />;
}

It actually works, but there is a problem, FancyInput only set current of ref accepted, not cleaning up.

Recall our explanation in React Internals Deep Dive 11 - how useRef() works?, React automatically cleans up refs attached to elements, but now it doesn’t.

What if ref changes during renders? Then the old ref would still hold the ref which causes problems since from the usage of <FancyInput ref={inputRef} />, it should be cleaned.

How to solve this ? We have useEffect() which could help clean things up, so we can try things like this.

jsx
function FancyInput(props, ref) {
const inputRef = useRef();
useEffect(() => {
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
return () => {
ref.current = null;
};
}, [ref]);
return <input ref={inputRef} />;
}
jsx
function FancyInput(props, ref) {
const inputRef = useRef();
useEffect(() => {
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
return () => {
ref.current = null;
};
}, [ref]);
return <input ref={inputRef} />;
}

but wait, how do you know for sure that ref is RefObject not function ref? Ok then, we need to check that.

js
function FancyInput(props, ref) {
const inputRef = useRef();
useEffect(() => {
if (typeof ref === "function") {
ref({
focus: () => {
inputRef.current.focus();
},
});
} else {
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
}
return () => {
if (typeof ref === "function") {
ref(null);
} else {
ref.current = null;
}
};
}, [ref]);
return <input ref={inputRef} />;
}
js
function FancyInput(props, ref) {
const inputRef = useRef();
useEffect(() => {
if (typeof ref === "function") {
ref({
focus: () => {
inputRef.current.focus();
},
});
} else {
ref.current = () => ({
focus: () => {
inputRef.current.focus();
},
});
}
return () => {
if (typeof ref === "function") {
ref(null);
} else {
ref.current = null;
}
};
}, [ref]);
return <input ref={inputRef} />;
}

You know what ? this is actually very similar to how useImperativeHandle() works. Except useImperativeHandle() is a layout effect, the ref setting happens at same stage of useLayoutEffect(), soonner than useEffect().

Ok let’s jump into the source code.

For effects, there is mount and update, which are different based on when useImperativeHandle() is called.

This is the simplified version of mountImperativeHandle(), (origin code)

js
function mountImperativeHandle<T>(
ref: {|current: T | null|} | ((inst: T | null) => mixed) | null | void,
create: () => T,
deps: Array<mixed> | void | null,
): void {
return mountEffectImpl(
fiberFlags,
HookLayout,
imperativeHandleEffect.bind(null, create, ref),
effectDeps,
);
}
js
function mountImperativeHandle<T>(
ref: {|current: T | null|} | ((inst: T | null) => mixed) | null | void,
create: () => T,
deps: Array<mixed> | void | null,
): void {
return mountEffectImpl(
fiberFlags,
HookLayout,
imperativeHandleEffect.bind(null, create, ref),
effectDeps,
);
}

Also for update, origin code

js
function updateImperativeHandle<T>(
ref: {| current: T | null |} | ((inst: T | null) => mixed) | null | void,
create: () => T,
deps: Array<mixed> | void | null
): void {
// TODO: If deps are provided, should we skip comparing the ref itself?
const effectDeps =
deps !== null && deps !== undefined ? deps.concat([ref]) : null;
return updateEffectImpl(
UpdateEffect,
HookLayout,
imperativeHandleEffect.bind(null, create, ref),
effectDeps
);
}
js
function updateImperativeHandle<T>(
ref: {| current: T | null |} | ((inst: T | null) => mixed) | null | void,
create: () => T,
deps: Array<mixed> | void | null
): void {
// TODO: If deps are provided, should we skip comparing the ref itself?
const effectDeps =
deps !== null && deps !== undefined ? deps.concat([ref]) : null;
return updateEffectImpl(
UpdateEffect,
HookLayout,
imperativeHandleEffect.bind(null, create, ref),
effectDeps
);
}

Notice that

  1. under the hood mountEffectImpl and updateEffectImpl are used. useEffect() and useLayoutEffect() does the same, here and here
  2. the 2nd argument is HookLayout, which means it is layout effect.

Last piece of puzzle, this is how imperativeHandleEffect() works. (code)

js
function imperativeHandleEffect<T>(
create: () => T,
ref: {| current: T | null |} | ((inst: T | null) => mixed) | null | void
) {
if (typeof ref === "function") {
const refCallback = ref;
const inst = create();
refCallback(inst);
return () => {
refCallback(null);
};
} else if (ref !== null && ref !== undefined) {
const refObject = ref;
const inst = create();
refObject.current = inst;
return () => {
refObject.current = null;
};
}
}
js
function imperativeHandleEffect<T>(
create: () => T,
ref: {| current: T | null |} | ((inst: T | null) => mixed) | null | void
) {
if (typeof ref === "function") {
const refCallback = ref;
const inst = create();
refCallback(inst);
return () => {
refCallback(null);
};
} else if (ref !== null && ref !== undefined) {
const refObject = ref;
const inst = create();
refObject.current = inst;
return () => {
refObject.current = null;
};
}
}

Set aside the details of perfectness, it actually looks very much the same as what we wrote, right?

Wrap-up

useImperativeHandle() is no magic, it just wraps the ref setting and cleaning for us, internally it is in the same stage as useLayoutEffect() so a bit sooner than useEffect().

Want to know more about how React works internally?
Check out my series - React Internals Deep Dive!

😳 Would you like to share my post to more people ?    

❮ Prev: How does useRef() work?

Next: Invariant-style assertion functions in TypeScript