feat: support createImmutable

src/Immutable.tsx

@@ -1,89 +1,106 @@
 import { supportRef } from 'rc-util/lib/ref';
 import * as React from 'react';
 
-const ImmutableContext = React.createContext<number>(0);
-
 export type CompareProps<T extends React.ComponentType<any>> = (
   prevProps: Readonly<React.ComponentProps<T>>,
   nextProps: Readonly<React.ComponentProps<T>>,
 ) => boolean;
 
 /**
- * Get render update mark by `makeImmutable` root.
- * Do not deps on the return value as render times
- * but only use for `useMemo` or `useCallback` deps.
+ * Create Immutable pair for `makeImmutable` and `responseImmutable`.
  */
-export function useImmutableMark() {
-  return React.useContext(ImmutableContext);
-}
+export default function createImmutable() {
+  const ImmutableContext = React.createContext<number>(null);
+
+  /**
+   * Get render update mark by `makeImmutable` root.
+   * Do not deps on the return value as render times
+   * but only use for `useMemo` or `useCallback` deps.
+   */
+  function useImmutableMark() {
+    return React.useContext(ImmutableContext);
+  }
 
-/**
+  /**
  * Wrapped Component will be marked as Immutable.
  * When Component parent trigger render,
  * it will notice children component (use with `responseImmutable`) node that parent has updated.
 
  * @param Component Passed Component
  * @param triggerRender Customize trigger `responseImmutable` children re-render logic. Default will always trigger re-render when this component re-render.
  */
-export function makeImmutable<T extends React.ComponentType<any>>(
-  Component: T,
-  shouldTriggerRender?: CompareProps<T>,
-): T {
-  const refAble = supportRef(Component);
-
-  const ImmutableComponent = function (props: any, ref: any) {
-    const refProps = refAble ? { ref } : {};
-    const renderTimesRef = React.useRef(0);
-    const prevProps = React.useRef(props);
-
-    if (
-      // Always trigger re-render if not provide `notTriggerRender`
-      !shouldTriggerRender ||
-      shouldTriggerRender(prevProps.current, props)
-    ) {
-      renderTimesRef.current += 1;
+  function makeImmutable<T extends React.ComponentType<any>>(
+    Component: T,
+    shouldTriggerRender?: CompareProps<T>,
+  ): T {
+    const refAble = supportRef(Component);
+
+    const ImmutableComponent = function (props: any, ref: any) {
+      const refProps = refAble ? { ref } : {};
+      const renderTimesRef = React.useRef(0);
+      const prevProps = React.useRef(props);
+
+      // If parent has the context, we do not wrap it
+      const mark = useImmutableMark();
+      if (mark !== null) {
+        return <Component {...props} {...refProps} />;
+      }
+
+      if (
+        // Always trigger re-render if not provide `notTriggerRender`
+        !shouldTriggerRender ||
+        shouldTriggerRender(prevProps.current, props)
+      ) {
+        renderTimesRef.current += 1;
+      }
+
+      prevProps.current = props;
+
+      return (
+        <ImmutableContext.Provider value={renderTimesRef.current}>
+          <Component {...props} {...refProps} />
+        </ImmutableContext.Provider>
+      );
+    };
+
+    if (process.env.NODE_ENV !== 'production') {
+      ImmutableComponent.displayName = `ImmutableRoot(${Component.displayName || Component.name})`;
     }
 
-    prevProps.current = props;
-
-    return (
-      <ImmutableContext.Provider value={renderTimesRef.current}>
-        <Component {...props} {...refProps} />
-      </ImmutableContext.Provider>
-    );
-  };
-
-  if (process.env.NODE_ENV !== 'production') {
-    ImmutableComponent.displayName = `ImmutableRoot(${Component.displayName || Component.name})`;
+    return refAble ? React.forwardRef(ImmutableComponent) : (ImmutableComponent as any);
   }
 
-  return refAble ? React.forwardRef(ImmutableComponent) : (ImmutableComponent as any);
-}
-
-/**
- * Wrapped Component with `React.memo`.
- * But will rerender when parent with `makeImmutable` rerender.
- */
-export function responseImmutable<T extends React.ComponentType<any>>(
-  Component: T,
-  propsAreEqual?: CompareProps<T>,
-): T {
-  const refAble = supportRef(Component);
-
-  const ImmutableComponent = function (props: any, ref: any) {
-    const refProps = refAble ? { ref } : {};
-    useImmutableMark();
-
-    return <Component {...props} {...refProps} />;
-  };
+  /**
+   * Wrapped Component with `React.memo`.
+   * But will rerender when parent with `makeImmutable` rerender.
+   */
+  function responseImmutable<T extends React.ComponentType<any>>(
+    Component: T,
+    propsAreEqual?: CompareProps<T>,
+  ): T {
+    const refAble = supportRef(Component);
+
+    const ImmutableComponent = function (props: any, ref: any) {
+      const refProps = refAble ? { ref } : {};
+      useImmutableMark();
+
+      return <Component {...props} {...refProps} />;
+    };
+
+    if (process.env.NODE_ENV !== 'production') {
+      ImmutableComponent.displayName = `ImmutableResponse(${
+        Component.displayName || Component.name
+      })`;
+    }
 
-  if (process.env.NODE_ENV !== 'production') {
-    ImmutableComponent.displayName = `ImmutableResponse(${
-      Component.displayName || Component.name
-    })`;
+    return refAble
+      ? React.memo(React.forwardRef(ImmutableComponent), propsAreEqual)
+      : (React.memo(ImmutableComponent, propsAreEqual) as any);
   }
 
-  return refAble
-    ? React.memo(React.forwardRef(ImmutableComponent), propsAreEqual)
-    : (React.memo(ImmutableComponent, propsAreEqual) as any);
+  return {
+    makeImmutable,
+    responseImmutable,
+    useImmutableMark,
+  };
 }

src/index.ts

@@ -1,6 +1,16 @@
 import type { SelectorContext } from './context';
 import { createContext, useContext } from './context';
-import { makeImmutable, responseImmutable, useImmutableMark } from './Immutable';
+import createImmutable from './Immutable';
 
-export { createContext, useContext, makeImmutable, responseImmutable, useImmutableMark };
+// For legacy usage, we export it directly
+const { makeImmutable, responseImmutable, useImmutableMark } = createImmutable();
+
+export {
+  createContext,
+  useContext,
+  createImmutable,
+  makeImmutable,
+  responseImmutable,
+  useImmutableMark,
+};
 export type { SelectorContext };

tests/immutable.test.tsx

@@ -1,6 +1,12 @@
 import { fireEvent, render } from '@testing-library/react';
 import React from 'react';
-import { createContext, makeImmutable, responseImmutable, useContext } from '../src';
+import {
+  createContext,
+  createImmutable,
+  makeImmutable,
+  responseImmutable,
+  useContext,
+} from '../src';
 import { RenderTimer, Value } from './common';
 
 describe('Immutable', () => {
@@ -165,4 +171,51 @@ describe('Immutable', () => {
     rerender(<ImmutableInput value="not-same" onChange={() => {}} />);
     expect(container.querySelector('#input').textContent).toEqual('2');
   });
+
+  describe('createImmutable', () => {
+    const { responseImmutable: responseCreatedImmutable, makeImmutable: makeCreatedImmutable } =
+      createImmutable();
+
+    it('nest should follow root', () => {
+      // child
+      const Little = responseCreatedImmutable(() => <RenderTimer id="little" />);
+
+      // parent
+      const Bamboo = makeCreatedImmutable(() => (
+        <>
+          <RenderTimer id="bamboo" />
+          <Little />
+        </>
+      ));
+
+      // root
+      const Light = makeCreatedImmutable(() => {
+        const [times, setTimes] = React.useState(0);
+
+        return (
+          <>
+            <button onClick={() => setTimes(i => i + 1)}>{times}</button>
+            <RenderTimer id="light" />
+            <Bamboo />
+          </>
+        );
+      });
+
+      const { container, rerender } = render(<Light />);
+
+      for (let i = 0; i < 10; i += 1) {
+        rerender(<Light />);
+      }
+      expect(container.querySelector('#light')!.textContent).toEqual('11');
+      expect(container.querySelector('#bamboo')!.textContent).toEqual('11');
+      expect(container.querySelector('#little')!.textContent).toEqual('11');
+
+      for (let i = 0; i < 10; i += 1) {
+        fireEvent.click(container.querySelector('button')!);
+      }
+      expect(container.querySelector('#light')!.textContent).toEqual('21');
+      expect(container.querySelector('#bamboo')!.textContent).toEqual('21');
+      expect(container.querySelector('#little')!.textContent).toEqual('11');
+    });
+  });
 });