fix case where we would omit the ref from reused vnodes (#3696)

* fix case where we would omit the ref from reused vnodes

* add test for call ordering

* correct tests

* skip test for now

src/diff/children.js

@@ -83,7 +83,7 @@ export function diffChildren(
 				childVNode.type,
 				childVNode.props,
 				childVNode.key,
-				null,
+				childVNode.ref ? childVNode.ref : null,
 				childVNode._original
 			);
 		} else {
@@ -244,14 +244,7 @@ function reorderChildren(childVNode, oldDom, parentDom) {
 			if (typeof vnode.type == 'function') {
 				oldDom = reorderChildren(vnode, oldDom, parentDom);
 			} else {
-				oldDom = placeChild(
-					parentDom,
-					vnode,
-					vnode,
-					c,
-					vnode._dom,
-					oldDom
-				);
+				oldDom = placeChild(parentDom, vnode, vnode, c, vnode._dom, oldDom);
 			}
 		}
 	}

src/diff/index.js

@@ -497,7 +497,9 @@ export function unmount(vnode, parentVNode, skipRemove) {
 	if (options.unmount) options.unmount(vnode);
 
 	if ((r = vnode.ref)) {
-		if (!r.current || r.current === vnode._dom) applyRef(r, null, parentVNode);
+		if (!r.current || r.current === vnode._dom) {
+			applyRef(r, null, parentVNode);
+		}
 	}
 
 	if ((r = vnode._component) != null) {

test/browser/refs.test.js

@@ -478,4 +478,69 @@ describe('refs', () => {
 		render(<App />, scratch);
 		expect(el).to.not.be.equal(null);
 	});
+
+	it('should not remove refs for memoized components keyed', () => {
+		const ref = createRef();
+		const element = <div ref={ref}>hey</div>;
+		function App(props) {
+			return <div key={props.count}>{element}</div>;
+		}
+
+		render(<App count={0} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+		render(<App count={1} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+		render(<App count={2} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+	});
+
+	it('should not remove refs for memoized components unkeyed', () => {
+		const ref = createRef();
+		const element = <div ref={ref}>hey</div>;
+		function App(props) {
+			return <div>{element}</div>;
+		}
+
+		render(<App count={0} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+		render(<App count={1} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+		render(<App count={2} />, scratch);
+		expect(ref.current).to.equal(scratch.firstChild.firstChild);
+	});
+
+	// TODO
+	it.skip('should properly call null for memoized components keyed', () => {
+		const calls = [];
+		const element = <div ref={x => calls.push(x)}>hey</div>;
+		function App(props) {
+			return <div key={props.count}>{element}</div>;
+		}
+
+		render(<App count={0} />, scratch);
+		render(<App count={1} />, scratch);
+		render(<App count={2} />, scratch);
+		expect(calls.length).to.equal(5);
+		expect(calls).to.deep.equal([
+			scratch.firstChild.firstChild,
+			null,
+			scratch.firstChild.firstChild,
+			null,
+			scratch.firstChild.firstChild
+		]);
+	});
+
+	it('should properly call null for memoized components unkeyed', () => {
+		const calls = [];
+		const element = <div ref={x => calls.push(x)}>hey</div>;
+		function App(props) {
+			return <div>{element}</div>;
+		}
+
+		render(<App count={0} />, scratch);
+		render(<App count={1} />, scratch);
+		render(<App count={2} />, scratch);
+		expect(calls.length).to.equal(1);
+		expect(calls[0]).to.equal(scratch.firstChild.firstChild);
+	});
 });