fix: ref logic

Author: zombieJ

src/Immutable.tsx

@@ -28,18 +28,20 @@ export function makeImmutable<T extends React.ComponentType<any>>(Component: T):
 export function responseImmutable<T extends React.ComponentType<any>>(Component: T): T {
   const refAble = supportRef(Component);
 
-  const ImmutableComponent = React.memo(function (props: any, ref: any) {
+  const ImmutableComponent = function (props: any, ref: any) {
     const refProps = refAble ? { ref } : {};
     React.useContext(RenderContext);
 
     return <Component {...props} {...refProps} />;
-  });
+  };
 
   if (process.env.NODE_ENV !== 'production') {
     ImmutableComponent.displayName = `ImmutableResponse(${
       Component.displayName || Component.name
     })`;
   }
 
-  return refAble ? React.forwardRef(ImmutableComponent) : (ImmutableComponent as any);
+  return refAble
+    ? React.memo(React.forwardRef(ImmutableComponent))
+    : (React.memo(ImmutableComponent) as any);
 }

tests/immutable.test.tsx

@@ -74,4 +74,25 @@ describe('Immutable', () => {
       expect(container.querySelector('#value')!.textContent).toEqual('1');
     });
   });
+
+  it('ref-able', () => {
+    const Root = React.forwardRef<HTMLDivElement, { children?: React.ReactNode }>((_, ref) => (
+      <div className="root" ref={ref} />
+    ));
+    const ImmutableRoot = makeImmutable(Root);
+    const Raw = React.forwardRef<HTMLDivElement>((_, ref) => <div className="raw" ref={ref} />);
+    const ImmutableRaw = responseImmutable(Raw);
+
+    const rootRef = React.createRef<HTMLDivElement>();
+    const rawRef = React.createRef<HTMLDivElement>();
+
+    const { container } = render(
+      <ImmutableRoot ref={rootRef}>
+        <ImmutableRaw ref={rawRef} />
+      </ImmutableRoot>,
+    );
+
+    expect(rootRef.current).toBe(container.querySelector('.root'));
+    expect(rawRef.current).toBe(container.querySelector('.raw'));
+  });
 });