fix(react-jss): Media query styles applied only to the first element in function

Author: pranshuchittora

packages/react-jss/.size-snapshot.json

@@ -1,23 +1,23 @@
 {
   "dist/react-jss.js": {
-    "bundled": 168964,
-    "minified": 58338,
-    "gzipped": 19097
+    "bundled": 170089,
+    "minified": 58766,
+    "gzipped": 19233
   },
   "dist/react-jss.min.js": {
-    "bundled": 111999,
-    "minified": 41597,
-    "gzipped": 14133
+    "bundled": 113124,
+    "minified": 42025,
+    "gzipped": 14254
   },
   "dist/react-jss.cjs.js": {
-    "bundled": 25873,
-    "minified": 11130,
-    "gzipped": 3722
+    "bundled": 26922,
+    "minified": 11621,
+    "gzipped": 3848
   },
   "dist/react-jss.esm.js": {
-    "bundled": 24911,
-    "minified": 10295,
-    "gzipped": 3600,
+    "bundled": 25951,
+    "minified": 10777,
+    "gzipped": 3726,
     "treeshaked": {
       "rollup": {
         "code": 1838,

packages/react-jss/src/createUseStyles.js

@@ -20,6 +20,13 @@ const useEffectOrLayoutEffect = isInBrowser ? React.useLayoutEffect : React.useE
 
 const noTheme = {}
 
+const reducer = (prevState, action) => {
+  if (action.type === 'updateSheet') {
+    return action.payload
+  }
+  return prevState
+}
+
 const createUseStyles = <Theme: {}>(styles: Styles<Theme>, options?: HookOptions<Theme> = {}) => {
   const {index = getSheetIndex(), theming, name, ...sheetOptions} = options
   const ThemeContext = (theming && theming.context) || DefaultThemeContext
@@ -35,75 +42,109 @@ const createUseStyles = <Theme: {}>(styles: Styles<Theme>, options?: HookOptions
     const context = React.useContext(JssContext)
     const theme = useTheme()
 
-    const [sheet, dynamicRules] = React.useMemo(
+    const [state, dispatch] = React.useReducer(reducer, null, () => {
+      const sheet = createStyleSheet({
+        context,
+        styles,
+        name,
+        theme,
+        index,
+        sheetOptions
+      })
+
+      let dynamicRules
+      let classes
+      if (sheet) {
+        if (context.registry) {
+          context.registry.add(sheet)
+        }
+        dynamicRules = addDynamicRules(sheet, data)
+        classes = getSheetClasses(sheet, dynamicRules)
+      }
+
+      return {
+        sheet,
+        dynamicRules,
+        classes: classes || {}
+      }
+    })
+    React.useMemo(
+      () => {
+        if (!isFirstMount.current) {
+          const newSheet = createStyleSheet({
+            context,
+            styles,
+            name,
+            theme,
+            index,
+            sheetOptions
+          })
+          const newDynamicRules = newSheet && addDynamicRules(newSheet, data)
+          const newClasses = newSheet ? getSheetClasses(newSheet, newDynamicRules) : {}
+
+          dispatch({
+            type: 'updateSheet',
+            payload: {
+              sheet: newSheet,
+              dynamicRules: newDynamicRules,
+              classes: newClasses
+            }
+          })
+        }
+      },
+      [theme, context]
+    )
+    useEffectOrLayoutEffect(
       () => {
-        const newSheet = createStyleSheet({
-          context,
-          styles,
-          name,
-          theme,
-          index,
-          sheetOptions
-        })
-
-        const newDynamicRules = newSheet ? addDynamicRules(newSheet, data) : null
-
-        if (newSheet) {
+        if (state.sheet) {
           manageSheet({
             index,
             context,
-            sheet: newSheet,
+            sheet: state.sheet,
             theme
           })
         }
 
-        return [newSheet, newDynamicRules]
+        return () => {
+          const {sheet, dynamicRules} = state
+
+          if (!sheet) return
+
+          unmanageSheet({
+            index,
+            context,
+            sheet,
+            theme
+          })
+
+          if (dynamicRules) {
+            removeDynamicRules(sheet, dynamicRules)
+          }
+        }
       },
-      [context, theme]
+      [state.sheet]
     )
 
     useEffectOrLayoutEffect(
       () => {
         // We only need to update the rules on a subsequent update and not in the first mount
-        if (sheet && dynamicRules && !isFirstMount.current) {
-          updateDynamicRules(data, sheet, dynamicRules)
+        if (state.sheet && state.dynamicRules && !isFirstMount.current) {
+          updateDynamicRules(data, state.sheet, state.dynamicRules)
         }
       },
       [data]
     )
 
-    useEffectOrLayoutEffect(
-      () =>
-        // cleanup only
-        () => {
-          if (sheet) {
-            unmanageSheet({
-              index,
-              context,
-              sheet,
-              theme
-            })
-          }
-
-          if (sheet && dynamicRules) {
-            removeDynamicRules(sheet, dynamicRules)
-          }
-        },
-      [sheet]
-    )
-
-    const classes = sheet && dynamicRules ? getSheetClasses(sheet, dynamicRules) : {}
-
     // $FlowFixMe
-    React.useDebugValue(classes)
+    React.useDebugValue(state.classes)
     // $FlowFixMe
     React.useDebugValue(theme === noTheme ? 'No theme' : theme)
 
     React.useEffect(() => {
       isFirstMount.current = false
     })
 
-    return classes
+    return state.classes
   }
 }