From effca36bb1c46313edb92756cb5940a15abe5bbd Mon Sep 17 00:00:00 2001
From: Pete Feltham <dev@felthy.com>
Date: Tue, 18 Jun 2019 19:45:50 +1000
Subject: [PATCH] Fix wrapped refs being called on every render

---
 src/waypoint.jsx       |  22 ++++-----
 test/waypoint_test.jsx | 101 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 111 insertions(+), 12 deletions(-)

diff --git a/src/waypoint.jsx b/src/waypoint.jsx
index addbfcd..c491676 100644
--- a/src/waypoint.jsx
+++ b/src/waypoint.jsx
@@ -36,6 +36,15 @@ export class Waypoint extends React.PureComponent {
 
     this.refElement = (e) => {
       this._ref = e;
+
+      const { children } = this.props;
+      if (children && children.ref && (isDOMElement(children) || isForwardRef(children))) {
+        if (typeof children.ref === 'function') {
+          children.ref(e);
+        } else {
+          children.ref.current = e;
+        }
+      }
     };
   }
 
@@ -302,18 +311,7 @@ export class Waypoint extends React.PureComponent {
     }
 
     if (isDOMElement(children) || isForwardRef(children)) {
-      const ref = (node) => {
-        this.refElement(node);
-        if (children.ref) {
-          if (typeof children.ref === 'function') {
-            children.ref(node);
-          } else {
-            children.ref.current = node;
-          }
-        }
-      };
-
-      return React.cloneElement(children, { ref });
+      return React.cloneElement(children, { ref: this.refElement });
     }
 
     return React.cloneElement(children, { innerRef: this.refElement });
diff --git a/test/waypoint_test.jsx b/test/waypoint_test.jsx
index 63ece51..2e66cb3 100644
--- a/test/waypoint_test.jsx
+++ b/test/waypoint_test.jsx
@@ -883,6 +883,22 @@ describe('<Waypoint>', () => {
       expect(this.subject).not.toThrow();
     });
 
+    it('does not throw with a Stateful Component that uses forwardRef as a child', () => {
+      class StatefulComponent extends React.Component {
+        render() {
+          const { forwardedRef } = this.props;
+          return <div ref={forwardedRef} />;
+        }
+      }
+
+      const ForwardedRefComponent = React.forwardRef((props, ref) => (
+        <StatefulComponent {...props} forwardedRef={ref} />
+      ));
+
+      this.props.children = <ForwardedRefComponent />;
+      expect(this.subject).not.toThrow();
+    });
+
     it('errors when a Stateful Component does not provide ref to Waypoint', () => {
       // eslint-disable-next-line react/prefer-stateless-function
       class StatefulComponent extends React.Component {
@@ -908,6 +924,91 @@ describe('<Waypoint>', () => {
       this.props.children = <StatelessComponent />;
       expect(this.subject).toThrowError(refNotUsedErrorMessage);
     });
+
+    it('wraps a DOM element’s existing ref', () => {
+      const functionRef = jasmine.createSpy('functionRef');
+      this.props.children = <div ref={functionRef} />;
+      expect(this.subject).not.toThrow();
+      expect(functionRef).toHaveBeenCalledTimes(1);
+      expect(functionRef).toHaveBeenCalledWith(jasmine.any(Node));
+    });
+
+    it('wraps a stateful component’s existing ref', () => {
+      const functionRef = jasmine.createSpy('functionRef');
+
+      const ForwardedRefComponent = React.forwardRef((props, ref) => (
+        <div ref={ref} />
+      ));
+
+      this.props.children = <ForwardedRefComponent ref={functionRef} />;
+      expect(this.subject).not.toThrow();
+      expect(functionRef).toHaveBeenCalledTimes(1);
+      expect(functionRef).toHaveBeenCalledWith(jasmine.any(Node));
+    });
+  });
+
+  describe('when the Waypoint child has its own ref', () => {
+    beforeEach(() => {
+      // Contrive so that onEnter triggers a re-render
+      class Wrapper extends React.Component {
+        render() {
+          const { onEnter, ...rest } = this.props;
+          const doOnEnter = () => {
+            onEnter();
+            this.forceUpdate();
+          };
+
+          return (
+            <Waypoint onEnter={doOnEnter} {...rest} />
+          );
+        }
+      }
+
+      this.subject = () => {
+        const el = renderAttached(
+          <div style={this.parentStyle}>
+            <div style={{ height: this.topSpacerHeight }} />
+            <Wrapper {...this.props} />
+            <div style={{ height: this.bottomSpacerHeight }} />
+          </div>,
+        );
+
+        jasmine.clock().tick(1);
+        return el;
+      };
+
+      this.topSpacerHeight = 200;
+      this.bottomSpacerHeight = 200;
+    });
+
+    it('does not cause wrapped DOM refs to be called on every render', () => {
+      const functionRef = jasmine.createSpy('functionRef');
+      this.props.children = <div style={{ height: 80 }} ref={functionRef} />;
+
+      const scrollable = this.subject();
+
+      scrollNodeTo(scrollable, 200);
+      scrollNodeTo(scrollable, 0);
+
+      expect(functionRef).toHaveBeenCalledTimes(1);
+    });
+
+    it('does not cause wrapped forwarded refs to be called on every render', () => {
+      const functionRef = jasmine.createSpy('functionRef');
+
+      const ForwardedRefComponent = React.forwardRef((props, ref) => (
+        <div style={{ height: 80 }} ref={ref} />
+      ));
+
+      this.props.children = <ForwardedRefComponent ref={functionRef} />;
+
+      const scrollable = this.subject();
+
+      scrollNodeTo(scrollable, 200);
+      scrollNodeTo(scrollable, 0);
+
+      expect(functionRef).toHaveBeenCalledTimes(1);
+    });
   });
 
   describe('when the Waypoint has children and is above the top', () => {