Add support for useLayoutEffect via exec.layout API

README.md

@@ -13,6 +13,7 @@ If you know how to [`useReducer`](https://reactjs.org/docs/hooks-reference.html#
 
 
 - [Installation](#installation)
+- [Isn't this unsafe?](#isnt-this-unsafe)
 - [Quick Start](#quick-start)
 - [Named Effects](#named-effects)
 - [Effect Implementations](#effect-implementations)

src/index.tsx

@@ -1,4 +1,11 @@
-import { useReducer, useEffect, useCallback, useRef, useMemo } from 'react';
+import {
+  useReducer,
+  useEffect,
+  useLayoutEffect,
+  useCallback,
+  useRef,
+  useMemo,
+} from 'react';
 
 type CleanupFunction = () => void;
 
@@ -18,6 +25,10 @@ export interface EffectObject<TState, TEvent extends EventObject> {
   exec?: EffectFunction<TState, TEvent, any>;
 }
 
+interface InternalEffectObject<TState, TEvent extends EventObject>
+  extends Record<typeof layoutEffectsSymbol, boolean>,
+    EffectObject<TState, TEvent> {}
+
 export type Effect<
   TState,
   TEvent extends EventObject,
@@ -32,6 +43,8 @@ type EntityTuple<TState, TEvent extends EventObject> = [
 type AggregatedEffectsState<TState, TEvent extends EventObject> = [
   TState,
   EntityTuple<TState, TEvent>[],
+  EffectEntity<TState, TEvent>[],
+  EntityTuple<TState, TEvent>[],
   EffectEntity<TState, TEvent>[]
 ];
 
@@ -46,7 +59,8 @@ enum EntityStatus {
   Stopped,
 }
 
-export interface EffectEntity<TState, TEvent extends EventObject> {
+export interface EffectEntity<TState, TEvent extends EventObject>
+  extends Record<typeof layoutEffectsSymbol, boolean> {
   type: string;
   status: EntityStatus;
   start: (state: TState, dispatch: React.Dispatch<TEvent>) => void;
@@ -56,13 +70,14 @@ export interface EffectEntity<TState, TEvent extends EventObject> {
 function createEffectEntity<
   TState,
   TEvent extends EventObject,
-  TEffect extends EffectObject<TState, TEvent>
+  TEffect extends InternalEffectObject<TState, TEvent>
 >(effect: TEffect): EffectEntity<TState, TEvent> {
   let effectCleanup: CleanupFunction | void;
 
   const entity: EffectEntity<TState, TEvent> = {
     type: effect.type,
     status: EntityStatus.Idle,
+    [layoutEffectsSymbol]: !!effect[layoutEffectsSymbol],
     start: (state, dispatch) => {
       if (effect.exec) {
         effectCleanup = effect.exec(state, effect, dispatch);
@@ -94,6 +109,9 @@ export interface EffectReducerExec<
     entity: EffectEntity<TState, TEvent> | undefined,
     effect: TEffect | EffectFunction<TState, TEvent, TEffect>
   ) => EffectEntity<TState, TEvent>;
+  layout: (
+    effect: TEffect | EffectFunction<TState, TEvent, TEffect>
+  ) => EffectEntity<TState, TEvent>;
 }
 
 export type EffectReducer<
@@ -106,6 +124,7 @@ export type EffectReducer<
   exec: EffectReducerExec<TState, TEvent, TEffect>
 ) => TState;
 
+const layoutEffectsSymbol = Symbol();
 const flushEffectsSymbol = Symbol();
 
 // 🚽
@@ -152,16 +171,26 @@ const toEffectObject = <
 >(
   effect: TEffect | EffectFunction<TState, TEvent, TEffect>,
   effectsMap?: EffectsMap<TState, TEvent, TEffect>
-): TEffect => {
+): InternalEffectObject<TState, TEvent> => {
   const type = typeof effect === 'function' ? effect.name : effect.type;
+  const layoutEffect =
+    typeof effect === 'function'
+      ? false
+      : // @ts-ignore
+        !!effect[layoutEffectsSymbol];
   const customExec = effectsMap
     ? effectsMap[type as TEffect['type']]
     : undefined;
   const exec =
     customExec || (typeof effect === 'function' ? effect : effect.exec);
   const other = typeof effect === 'function' ? {} : effect;
 
-  return { ...other, type, exec } as TEffect;
+  return {
+    ...other,
+    type,
+    exec,
+    [layoutEffectsSymbol]: layoutEffect,
+  } as InternalEffectObject<TState, TEvent>;
 };
 
 export type InitialEffectStateGetter<
@@ -184,47 +213,84 @@ export function useEffectReducer<
   effectsMap?: EffectsMap<TState, TEvent, TEffect>
 ): [TState, React.Dispatch<TEvent | TEvent['type']>] {
   const entitiesRef = useRef<Set<EffectEntity<TState, TEvent>>>(new Set());
+  const layoutEntitiesRef = useRef<Set<EffectEntity<TState, TEvent>>>(
+    new Set()
+  );
   const wrappedReducer = (
-    [state, stateEffectTuples, entitiesToStop]: AggregatedEffectsState<
-      TState,
-      TEvent
-    >,
+    [
+      state,
+      effectStateTuples,
+      entitiesToStop,
+      layoutEffectStateTuples,
+      layoutEntitiesToStop,
+    ]: AggregatedEffectsState<TState, TEvent>,
     event: TEvent | FlushEvent
   ): AggregatedEffectsState<TState, TEvent> => {
     const nextEffectEntities: Array<EffectEntity<TState, TEvent>> = [];
     const nextEntitiesToStop: Array<EffectEntity<TState, TEvent>> = [];
+    const nextLayoutEffectEntities: Array<EffectEntity<TState, TEvent>> = [];
+    const nextLayoutEntitiesToStop: Array<EffectEntity<TState, TEvent>> = [];
 
     if (event.type === flushEffectsSymbol) {
       // Record that effects have already been executed
-      return [state, stateEffectTuples.slice(event.count), nextEntitiesToStop];
+      return [
+        state,
+        effectStateTuples.slice(event.count),
+        nextEntitiesToStop,
+        layoutEffectStateTuples.slice(event.count),
+        nextLayoutEntitiesToStop,
+      ];
     }
 
     const exec = (
       effect: TEffect | EffectFunction<TState, TEvent, TEffect>
     ) => {
       const effectObject = toEffectObject(effect, effectsMap);
-      const effectEntity = createEffectEntity<TState, TEvent, TEffect>(
-        effectObject
-      );
-      nextEffectEntities.push(effectEntity);
+      const effectEntity = createEffectEntity<
+        TState,
+        TEvent,
+        InternalEffectObject<TState, TEvent>
+      >(effectObject);
+
+      (effectObject[layoutEffectsSymbol]
+        ? nextLayoutEffectEntities
+        : nextEffectEntities
+      ).push(effectEntity);
 
       return effectEntity;
     };
 
     exec.stop = (entity: EffectEntity<TState, TEvent>) => {
-      nextEntitiesToStop.push(entity);
+      (entity[layoutEffectsSymbol]
+        ? nextLayoutEntitiesToStop
+        : nextEntitiesToStop
+      ).push(entity);
     };
 
     exec.replace = (
       entity: EffectEntity<TState, TEvent>,
       effect: TEffect | EffectFunction<TState, TEvent, TEffect>
     ) => {
       if (entity) {
-        nextEntitiesToStop.push(entity);
+        (entity[layoutEffectsSymbol]
+          ? nextLayoutEntitiesToStop
+          : nextEntitiesToStop
+        ).push(entity);
       }
       return exec(effect);
     };
 
+    exec.layout = (
+      effect: TEffect | EffectFunction<TState, TEvent, TEffect>
+    ) => {
+      const effectObject = toEffectObject(effect, effectsMap);
+      return exec({
+        ...effectObject,
+        // @ts-ignore
+        [layoutEffectsSymbol]: true,
+      });
+    };
+
     const nextState = effectReducer(
       state,
       event,
@@ -234,11 +300,17 @@ export function useEffectReducer<
     return [
       nextState,
       nextEffectEntities.length
-        ? [...stateEffectTuples, [nextState, nextEffectEntities]]
-        : stateEffectTuples,
+        ? [...effectStateTuples, [nextState, nextEffectEntities]]
+        : effectStateTuples,
       entitiesToStop.length
         ? [...entitiesToStop, ...nextEntitiesToStop]
         : nextEntitiesToStop,
+      nextLayoutEffectEntities.length
+        ? [...layoutEffectStateTuples, [nextState, nextLayoutEffectEntities]]
+        : layoutEffectStateTuples,
+      layoutEntitiesToStop.length
+        ? [...layoutEntitiesToStop, ...nextLayoutEntitiesToStop]
+        : nextLayoutEntitiesToStop,
     ];
   };
 
@@ -248,33 +320,50 @@ export function useEffectReducer<
   > = useMemo(() => {
     if (typeof initialState === 'function') {
       const initialEffectEntities: Array<EffectEntity<TState, TEvent>> = [];
+      const initialLayoutEffectEntities: Array<EffectEntity<
+        TState,
+        TEvent
+      >> = [];
 
       const resolvedInitialState = (initialState as InitialEffectStateGetter<
         TState,
         TEvent,
         TEffect
       >)(effect => {
         const effectObject = toEffectObject(effect, effectsMap);
-        const effectEntity = createEffectEntity<TState, TEvent, TEffect>(
-          effectObject
-        );
-
-        initialEffectEntities.push(effectEntity);
+        const effectEntity = createEffectEntity<
+          TState,
+          TEvent,
+          InternalEffectObject<TState, TEvent>
+        >(effectObject);
+
+        (effectObject[layoutEffectsSymbol]
+          ? initialLayoutEffectEntities
+          : initialEffectEntities
+        ).push(effectEntity);
         return effectEntity;
       });
 
       return [
         resolvedInitialState,
         [[resolvedInitialState, initialEffectEntities]],
         [],
+        [[resolvedInitialState, initialLayoutEffectEntities]],
+        [],
       ];
     }
 
-    return [initialState, [], []];
+    return [initialState, [], [], [], []];
   }, []);
 
   const [
-    [state, effectStateEntityTuples, entitiesToStop],
+    [
+      state,
+      effectStateEntityTuples,
+      entitiesToStop,
+      layoutEffectStateEntityTuples,
+      layoutEntitiesToStop,
+    ],
     dispatch,
   ] = useReducer(wrappedReducer, initialStateAndEffects);
 
@@ -325,5 +414,43 @@ export function useEffectReducer<
     };
   }, []);
 
+  // Now do it all again, but with layout effects
+  useLayoutEffect(() => {
+    if (layoutEntitiesToStop.length) {
+      layoutEntitiesToStop.forEach(entity => {
+        entity.stop();
+        layoutEntitiesRef.current.delete(entity);
+      });
+    }
+  }, [layoutEntitiesToStop]);
+
+  useLayoutEffect(() => {
+    if (layoutEffectStateEntityTuples.length) {
+      layoutEffectStateEntityTuples.forEach(([effectState, effectEntities]) => {
+        effectEntities.forEach(entity => {
+          if (entity.status !== EntityStatus.Idle) return;
+
+          layoutEntitiesRef.current.add(entity);
+          entity.start(effectState, dispatch);
+        });
+      });
+
+      dispatch({
+        type: flushEffectsSymbol,
+        count: layoutEffectStateEntityTuples.length,
+      });
+    }
+  }, [layoutEffectStateEntityTuples]);
+
+  useLayoutEffect(() => {
+    return () => {
+      layoutEntitiesRef.current.forEach(entity => {
+        if (entity.status === EntityStatus.Started) {
+          entity.stop();
+        }
+      });
+    };
+  }, []);
+
   return [state, wrappedDispatch];
 }