Render lazy-loaded components synchronously if already loaded

Make lazy-loading components created by `lazy` render synchronously (ie. render
the actual component on the first render) if the component was already loaded by
a previous instance of the lazy wrapper component.

In h the async rendering was causing annotation cards to transition their height
slightly after appearing. This change doesn't prevent the animation from
happening when annotation cards first appear, but does prevent it on subsequen
renders.

Author: robertknight

src/util/lazy.tsx

@@ -43,8 +43,14 @@ export function lazy<P>(
   load: () => Promise<ComponentModule<P> | FunctionalComponent<P>>,
   { errorFallback, fallback }: LazyOptions<P>,
 ) {
+  // Cache resolved component so that instances created after it has loaded
+  // will render synchronously.
+  let resolved: FunctionalComponent<P> | undefined;
+
   function Lazy(props: P & JSX.IntrinsicAttributes) {
-    const [component, setComponent] = useState<FunctionalComponent<P>>();
+    const [component, setComponent] = useState<
+      FunctionalComponent<P> | undefined
+    >(() => resolved);
     const [error, setError] = useState<Error | null>(null);
     const [loading, setLoading] = useState(false);
 
@@ -63,12 +69,18 @@ export function lazy<P>(
     if (!component && !loading) {
       setLoading(true);
       load()
-        .then(component => {
-          if (isComponentModule(component)) {
-            setComponent(() => component.default);
+        .then(componentOrModule => {
+          let component;
+          if (isComponentModule(componentOrModule)) {
+            component = componentOrModule.default;
           } else {
-            setComponent(() => component);
+            component = componentOrModule as FunctionalComponent<P>;
           }
+
+          // Cache so we can render synchronously in future.
+          resolved = component;
+
+          setComponent(() => component);
         })
         .catch(setError)
         .finally(() => {

src/util/test/lazy-test.js

@@ -55,6 +55,21 @@ describe('lazy', () => {
     );
   });
 
+  it('renders component on initial render if already loaded', async () => {
+    // First render, which triggers lazy loading.
+    fakeLoader.returns(Promise.resolve(fakeComponent));
+    const wrapper = mount(<LazyComponent text="test" />);
+    await delay(0);
+    wrapper.update();
+    assert.isTrue(wrapper.exists('[data-testid="loaded-component"]'));
+    wrapper.unmount();
+
+    // Second render, which should synchronously render the already-loaded
+    // component.
+    const wrapper2 = mount(<LazyComponent text="test" />);
+    assert.isTrue(wrapper2.exists('[data-testid="loaded-component"]'));
+  });
+
   it('supports load callback returning a module', async () => {
     fakeLoader.returns(Promise.resolve({ default: fakeComponent }));
     const wrapper = mount(<LazyComponent text="test" />);