Fix passing ref to functional components (#153)

src/index.tsx

@@ -1,4 +1,4 @@
-import React, { HTMLAttributes } from 'react';
+import React, { HTMLAttributes, RefAttributes, ReactInstance } from 'react';
 import ReactDOM from 'react-dom';
 import contains from 'rc-util/lib/Dom/contains';
 import findDOMNode from 'rc-util/lib/Dom/findDOMNode';
@@ -732,9 +732,8 @@ export function generateTrigger(
       const { popupVisible } = this.state;
       const { children, forceRender, alignPoint, className } = this.props;
       const child = React.Children.only(children) as React.ReactElement;
-      const newChildProps: HTMLAttributes<HTMLElement> & { key: string } = {
-        key: 'trigger',
-      };
+      const newChildProps: HTMLAttributes<HTMLElement> &
+        RefAttributes<ReactInstance> & { key: string } = { key: 'trigger' };
 
       if (this.isContextMenuToShow()) {
         newChildProps.onContextMenu = this.onContextMenu;
@@ -779,10 +778,17 @@ export function generateTrigger(
       if (childrenClassName) {
         newChildProps.className = childrenClassName;
       }
-      const trigger = React.cloneElement(child, {
-        ...newChildProps,
-        ref: composeRef(this.triggerRef, (child as any).ref),
-      });
+
+      // Prevent adding ref to Functional child components
+      if (
+        !child.type ||
+        typeof child.type !== 'function' ||
+        (child.type.prototype && child.type.prototype.isReactComponent)
+      ) {
+        newChildProps.ref = composeRef(this.triggerRef, (child as any).ref);
+      }
+
+      const trigger = React.cloneElement(child, newChildProps);
 
       let portal: React.ReactElement;
       // prevent unmounting after it's rendered

tests/basic.test.jsx

@@ -538,22 +538,19 @@ describe('Trigger.Basic', () => {
   });
 
   it('support function component', () => {
-    const NoRef = React.forwardRef((props, ref) => {
-      React.useImperativeHandle(ref, () => null);
-      return (
-        <div className="target" {...props}>
-          click
-        </div>
-      );
-    });
+    const FuncComp = props => (
+      <div className="target" {...props}>
+        click
+      </div>
+    );
 
     const wrapper = mount(
       <Trigger
         action={['click']}
         popupAlign={placementAlignMap.left}
         popup={<strong className="x-content">tooltip2</strong>}
       >
-        <NoRef />
+        <FuncComp />
       </Trigger>,
     );
 
@@ -564,6 +561,66 @@ describe('Trigger.Basic', () => {
     expect(wrapper.isHidden()).toBeTruthy();
   });
 
+  describe('passes ref to children where applicable', () => {
+    function getRefResult(component) {
+      const triggerMock = jest.fn();
+      const wrapper = mount(
+        <Trigger
+          action={['click']}
+          popupAlign={placementAlignMap.left}
+          popup={<strong className="x-content">tooltip2</strong>}
+          ref={triggerMock}
+        >
+          {component}
+        </Trigger>,
+      );
+
+      wrapper.trigger();
+
+      return triggerMock.mock.calls[0][0].triggerRef.current;
+    }
+
+    it('does not pass ref to function component', () => {
+      const NoRef = props => (
+        <div id="target" {...props}>
+          click
+        </div>
+      );
+
+      const refVal = getRefResult(<NoRef />);
+      expect(refVal).toBeNull();
+    });
+
+    it('does pass ref to forwardRef function component', () => {
+      const WithRef = React.forwardRef((props, ref) => (
+        <div id="target" {...props} ref={ref}>
+          click
+        </div>
+      ));
+
+      const refVal = getRefResult(<WithRef />);
+      expect(refVal.id).toBe('target');
+    });
+
+    it('does pass ref to class component', () => {
+      class ClassComp extends React.Component {
+        id = 'target';
+
+        render() {
+          return <div>click</div>;
+        }
+      }
+
+      const refVal = getRefResult(<ClassComp />);
+      expect(refVal.id).toBe('target');
+    });
+
+    it('does pass ref to element', () => {
+      const refVal = getRefResult(<div id="target">click</div>);
+      expect(refVal.id).toBe('target');
+    });
+  });
+
   it('Popup with mouseDown prevent', () => {
     const wrapper = mount(
       <Trigger