perf: Apply state colocation principles and avoid prop drilling

In particular the story context state is maintained near the root of the tree, any changes to this state end up re-rendering a large number of components, such as the story list. By reading the state closer to where it needs to be used, it is possible to avoid prop drilling-inflicted rerenders. I opted for Jotai here, since its a relatively small and well tested and maintained state management library, that imposes basically nothing on an application's existing structure.
storybookjs · Feb 16, 2023 · d877d29 · d877d29
1 parent a04e092
commit d877d29
Show file tree

Hide file tree

Showing 7 changed files with 98 additions and 50 deletions.
diff --git a/app/react-native/package.json b/app/react-native/package.json
@@ -63,6 +63,7 @@
     "commander": "^8.2.0",
     "emotion-theming": "^10.0.19",
     "glob": "^7.1.7",
+    "jotai": "^2.0.2",
     "prettier": "^2.4.1",
     "react-native-swipe-gestures": "^1.0.5",
     "util": "^0.12.4"

diff --git a/app/react-native/src/hooks.tsx b/app/react-native/src/hooks.tsx
@@ -0,0 +1,44 @@
+import { useMemo } from 'react';
+import type { StoryContext } from '@storybook/csf';
+import { atom, useAtomValue, useSetAtom } from 'jotai';
+
+import type { ReactNativeFramework } from './types/types-6.0';
+
+const storyContextAtom = atom(null as (StoryContext<ReactNativeFramework> | null));
+
+/**
+ * Hook that returns a function to set the current story context.
+ */
+export function useSetStoryContext() {
+  return useSetAtom(storyContextAtom);
+}
+
+/**
+ * Hook to read the current story context.
+ */
+export function useStoryContext() {
+  return useAtomValue(storyContextAtom);
+}
+
+export function useStoryContextParam<T = any>(name: string, defaultValue?: T): T {
+  const paramAtom = useMemo(
+    () => atom(get => get(storyContextAtom)?.parameters?.[name] ?? defaultValue),
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [name]
+  );
+  return useAtomValue(paramAtom);
+}
+
+/**
+ * Hook that indicates if `storyId` is the currently selected story.
+ */
+export function useIsStorySelected(storyId: string) {
+  return useAtomValue(useMemo(() => atom(get => get(storyContextAtom)?.id === storyId), [storyId]));
+}
+
+/**
+ * Hook that indicates if `title` is the currently selected story section.
+ */
+export function useIsStorySectionSelected(title: string) {
+  return useAtomValue(useMemo(() => atom(get => get(storyContextAtom)?.title === title), [title]));
+}
diff --git a/app/react-native/src/preview/View.tsx b/app/react-native/src/preview/View.tsx
@@ -1,10 +1,11 @@
-import React, { useEffect, useState, useReducer } from 'react';
+import React, { useEffect, useReducer } from 'react';
 import AsyncStorage from '@react-native-async-storage/async-storage';
 import { StoryIndex, SelectionSpecifier } from '@storybook/store';
 import { StoryContext, toId } from '@storybook/csf';
 import { addons } from '@storybook/addons';
 import { ThemeProvider } from 'emotion-theming';
 import { SafeAreaProvider } from 'react-native-safe-area-context';
+import { useSetStoryContext } from '../hooks';
 import OnDeviceUI from './components/OnDeviceUI';
 import { theme } from './components/Shared/theme';
 import type { ReactNativeFramework } from '../types/types-6.0';
@@ -143,7 +144,7 @@ export class View {
 
     const appliedTheme = { ...theme, ...params.theme };
     return () => {
-      const [context, setContext] = useState<StoryContext<ReactNativeFramework>>();
+      const setContext = useSetStoryContext();
       const [, forceUpdate] = useReducer((x) => x + 1, 0);
       useEffect(() => {
         self._setStory = (newStory: StoryContext<ReactNativeFramework>) => {
@@ -166,7 +167,6 @@ export class View {
           <SafeAreaProvider>
             <ThemeProvider theme={appliedTheme}>
               <OnDeviceUI
-                context={context}
                 storyIndex={self._storyIndex}
                 isUIHidden={params.isUIHidden}
                 tabOpen={params.tabOpen}
@@ -177,7 +177,7 @@ export class View {
           </SafeAreaProvider>
         );
       } else {
-        return <StoryView context={context} />;
+        return <StoryView />;
       }
     };
   };

diff --git a/app/react-native/src/preview/components/OnDeviceUI/OnDeviceUI.tsx b/app/react-native/src/preview/components/OnDeviceUI/OnDeviceUI.tsx
@@ -14,6 +14,7 @@ import {
   StyleSheet,
   View,
 } from 'react-native';
+import { useStoryContextParam } from '../../../hooks';
 import StoryListView from '../StoryListView';
 import StoryView from '../StoryView';
 import AbsolutePositionedKeyboardAwareView, {
@@ -32,8 +33,6 @@ import { PREVIEW, ADDONS } from './navigation/constants';
 import Panel from './Panel';
 import { useWindowDimensions } from 'react-native';
 import { useSafeAreaInsets } from 'react-native-safe-area-context';
-import { StoryContext } from '@storybook/csf';
-import { ReactNativeFramework } from 'src/types/types-6.0';
 
 const ANIMATION_DURATION = 300;
 const IS_IOS = Platform.OS === 'ios';
@@ -43,7 +42,6 @@ export const IS_EXPO = getExpoRoot() !== undefined;
 const IS_ANDROID = Platform.OS === 'android';
 const BREAKPOINT = 1024;
 interface OnDeviceUIProps {
-  context: StoryContext<ReactNativeFramework>;
   storyIndex: StoryIndex;
   url?: string;
   tabOpen?: number;
@@ -75,7 +73,6 @@ const styles = StyleSheet.create({
 });
 
 const OnDeviceUI = ({
-  context,
   storyIndex,
   isUIHidden,
   shouldDisableKeyboardAvoidingView,
@@ -112,7 +109,7 @@ const OnDeviceUI = ({
     }
   }, [tabOpen]);
 
-  const noSafeArea = context?.parameters?.noSafeArea ?? false;
+  const noSafeArea = useStoryContextParam<boolean>('noSafeArea', false);
   const previewWrapperStyles = [
     flex,
     getPreviewPosition({
@@ -146,7 +143,7 @@ const OnDeviceUI = ({
               <Animated.View style={previewStyles}>
                 <Preview disabled={tabOpen === PREVIEW}>
                   <WrapperView style={[flex, wrapperMargin]}>
-                    <StoryView context={context} />
+                    <StoryView />
                   </WrapperView>
                 </Preview>
                 {tabOpen !== PREVIEW ? (
@@ -164,7 +161,7 @@ const OnDeviceUI = ({
                 wide
               )}
             >
-              <StoryListView storyIndex={storyIndex} selectedStoryContext={context} />
+              <StoryListView storyIndex={storyIndex} />
             </Panel>
 
             <Panel

diff --git a/app/react-native/src/preview/components/StoryListView/StoryListView.tsx b/app/react-native/src/preview/components/StoryListView/StoryListView.tsx
@@ -2,13 +2,17 @@ import styled from '@emotion/native';
 import { addons, StoryKind } from '@storybook/addons';
 import { StoryIndex, StoryIndexEntry } from '@storybook/client-api';
 import Events from '@storybook/core-events';
-import { StoryContext } from '@storybook/csf';
 import React, { useMemo, useState } from 'react';
-import { SectionList, StyleSheet, View } from 'react-native';
+import {
+  SectionList,
+  SectionListRenderItem,
+  StyleSheet,
+  View,
+} from 'react-native';
 import { useSafeAreaInsets } from 'react-native-safe-area-context';
 import { GridIcon, StoryIcon } from '../Shared/icons';
 import { Header, Name } from '../Shared/text';
-import { ReactNativeFramework } from 'src/types/types-6.0';
+import { useIsStorySelected, useIsStorySectionSelected } from '../../../hooks';
 
 const SearchBar = styled.TextInput(
   {
@@ -54,20 +58,22 @@ const StoryListContainer = styled.View(({ theme }) => ({
 
 interface SectionProps {
   title: string;
-  selected: boolean;
 }
 
-const SectionHeader = React.memo(({ title, selected }: SectionProps) => (
-  <HeaderContainer key={title}>
-    <GridIcon />
-    <Header selected={selected}>{title}</Header>
-  </HeaderContainer>
-));
+const SectionHeader = React.memo(({ title }: SectionProps) => {
+  const selected = useIsStorySectionSelected(title);
+  return (
+    <HeaderContainer key={title}>
+      <GridIcon />
+      <Header selected={selected}>{title}</Header>
+    </HeaderContainer>
+  );
+});
 
 interface ListItemProps {
+  storyId: string;
   title: string;
   kind: string;
-  selected: boolean;
   onPress: () => void;
 }
 
@@ -82,25 +88,27 @@ const ItemTouchable = styled.TouchableOpacity<{ selected: boolean }>(
 );
 
 const ListItem = React.memo(
-  ({ kind, title, selected, onPress }: ListItemProps) => (
-    <ItemTouchable
-      key={title}
-      onPress={onPress}
-      activeOpacity={0.8}
-      testID={`Storybook.ListItem.${kind}.${title}`}
-      accessibilityLabel={`Storybook.ListItem.${title}`}
-      selected={selected}
-    >
-      <StoryIcon selected={selected} />
-      <Name selected={selected}>{title}</Name>
-    </ItemTouchable>
-  ),
-  (prevProps, nextProps) => prevProps.selected === nextProps.selected
+  ({ storyId, kind, title, onPress }: ListItemProps) => {
+    const selected = useIsStorySelected(storyId);
+    return (
+      <ItemTouchable
+        key={title}
+        onPress={onPress}
+        activeOpacity={0.8}
+        testID={`Storybook.ListItem.${kind}.${title}`}
+        accessibilityLabel={`Storybook.ListItem.${title}`}
+        selected={selected}
+      >
+        <StoryIcon selected={selected} />
+        <Name selected={selected}>{title}</Name>
+      </ItemTouchable>
+    );
+  },
+  (prevProps, nextProps) => prevProps.storyId === nextProps.storyId
 );
 
 interface Props {
   storyIndex: StoryIndex;
-  selectedStoryContext: StoryContext<ReactNativeFramework>;
 }
 
 interface DataItem {
@@ -130,11 +138,11 @@ const styles = StyleSheet.create({
 
 const tabBarHeight = 40;
 
-const StoryListView = ({ selectedStoryContext, storyIndex }: Props) => {
 function keyExtractor(item: any, index) {
   return item.id + index;
 }
 
+const StoryListView = ({ storyIndex }: Props) => {
   const insets = useSafeAreaInsets();
   const originalData = useMemo(() => getStories(storyIndex), [storyIndex]);
   const [data, setData] = useState<DataItem[]>(originalData);
@@ -178,19 +186,16 @@ function keyExtractor(item: any, index) {
   const renderItem: SectionListRenderItem<StoryIndexEntry, DataItem> = React.useCallback(({item}) => {
     return (
       <ListItem
+        storyId={item.id}
         title={item.name}
         kind={item.title}
-        selected={selectedStoryContext && item.id === selectedStoryContext.id}
         onPress={() => changeStory(item.id)}
       />
     );
   }, []);
 
   const renderSectionHeader = React.useCallback(({ section: { title } }) => (
-    <SectionHeader
-      title={title}
-      selected={selectedStoryContext && title === selectedStoryContext.title}
-    />
+    <SectionHeader title={title} />
   ), []);
 
   return (

diff --git a/app/react-native/src/preview/components/StoryView/StoryView.tsx b/app/react-native/src/preview/components/StoryView/StoryView.tsx
@@ -1,12 +1,7 @@
 import React from 'react';
 
 import { Text, View, StyleSheet } from 'react-native';
-import type { StoryContext } from '@storybook/csf';
-import { ReactNativeFramework } from 'src/types/types-6.0';
-
-interface Props {
-  context?: StoryContext<ReactNativeFramework>;
-}
+import { useStoryContext } from '../../../hooks';
 
 const styles = StyleSheet.create({
   container: { flex: 1 },
@@ -19,7 +14,8 @@ const styles = StyleSheet.create({
   },
 });
 
-const StoryView = ({ context }: Props) => {
+const StoryView = () => {
+  const context = useStoryContext();
   const id = context?.id;
 
   if (context && context.unboundStoryFn) {

diff --git a/yarn.lock b/yarn.lock
@@ -13025,6 +13025,11 @@ joi@^17.2.1:
     "@sideway/formula" "^3.0.0"
     "@sideway/pinpoint" "^2.0.0"
 
+jotai@^2.0.2:
+  version "2.0.2"
+  resolved "https://registry.yarnpkg.com/jotai/-/jotai-2.0.2.tgz#52728e91de61742263e40fda79a367ab4649506c"
+  integrity sha512-0yOked08Swa84LUbBjtj7ZLZrE05n3u50rHeZ+bsT86thUjcy0kFgQz9GmEWhYbQDFoT1G8Ww6edSj/MBJHO4A==
+
 js-string-escape@^1.0.1:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/js-string-escape/-/js-string-escape-1.0.1.tgz#e2625badbc0d67c7533e9edc1068c587ae4137ef"