feat: Add option to render concurrent roots

package.json

@@ -55,8 +55,8 @@
     "cross-env": "^6.0.0",
     "kcd-scripts": "^1.7.0",
     "npm-run-all": "^4.1.5",
-    "react": "^16.9.0",
-    "react-dom": "^16.9.0",
+    "react": "^0.0.0-experimental-f6b8d31a7",
+    "react-dom": "^0.0.0-experimental-f6b8d31a7",
     "rimraf": "^3.0.0"
   },
   "peerDependencies": {

src/__tests__/concurrent.js

@@ -0,0 +1,52 @@
+import React from 'react'
+import {act, render, cleanup} from '../'
+
+test('simple render works like legacy', () => {
+  const {container} = render(<div>test</div>, {root: 'concurrent'})
+
+  expect(container).toHaveTextContent('test')
+})
+
+test('unmounts are flushed in sync', () => {
+  const {container, unmount} = render(<div>test</div>, {root: 'concurrent'})
+
+  unmount()
+
+  expect(container.children).toHaveLength(0)
+})
+
+test('rerender are flushed in sync', () => {
+  const {container, rerender} = render(<div>foo</div>, {root: 'concurrent'})
+
+  rerender(<div>bar</div>)
+
+  expect(container).toHaveTextContent('foo')
+})
+
+test('cleanup unmounts in sync', () => {
+  const {container} = render(<div>test</div>, {root: 'concurrent'})
+
+  cleanup()
+
+  expect(container.children).toHaveLength(0)
+})
+
+test('state updates are concurrent', () => {
+  function TrackingButton() {
+    const [clickCount, increment] = React.useReducer(n => n + 1, 0)
+
+    return (
+      <button type="button" onClick={increment}>
+        Clicked {clickCount} times.
+      </button>
+    )
+  }
+  const {getByRole} = render(<TrackingButton />, {root: 'concurrent'})
+
+  act(() => {
+    getByRole('button').click()
+    expect(getByRole('button')).toHaveTextContent('Clicked 0 times')
+  })
+
+  expect(getByRole('button')).toHaveTextContent('Clicked 1 times')
+})

src/pure.js

@@ -19,6 +19,10 @@ configureDTL({
 })
 
 const mountedContainers = new Set()
+/**
+ * @type {WeakMap<Element, ReactRoot>}
+ */
+const mountedRoots = new WeakMap()
 
 function render(
   ui,
@@ -28,6 +32,7 @@ function render(
     queries,
     hydrate = false,
     wrapper: WrapperComponent,
+    root,
   } = {},
 ) {
   if (!baseElement) {
@@ -44,13 +49,21 @@ function render(
   // they're passing us a custom container or not.
   mountedContainers.add(container)
 
+  let reactRoot = null
+  if (root === 'concurrent') {
+    reactRoot = ReactDOM.createRoot(container, {hydrate})
+    mountedRoots.set(container, reactRoot)
+  }
+
   const wrapUiIfNeeded = innerElement =>
     WrapperComponent
       ? React.createElement(WrapperComponent, null, innerElement)
       : innerElement
 
   act(() => {
-    if (hydrate) {
+    if (reactRoot) {
+      reactRoot.render(wrapUiIfNeeded(ui))
+    } else if (hydrate) {
       ReactDOM.hydrate(wrapUiIfNeeded(ui), container)
     } else {
       ReactDOM.render(wrapUiIfNeeded(ui), container)
@@ -66,11 +79,23 @@ function render(
           el.forEach(e => console.log(prettyDOM(e)))
         : // eslint-disable-next-line no-console,
           console.log(prettyDOM(el)),
-    unmount: () => ReactDOM.unmountComponentAtNode(container),
+    unmount: () => {
+      if (reactRoot) {
+        act(() => {
+          reactRoot.unmount()
+        })
+      } else {
+        ReactDOM.unmountComponentAtNode(container)
+      }
+    },
     rerender: rerenderUi => {
-      render(wrapUiIfNeeded(rerenderUi), {container, baseElement})
-      // Intentionally do not return anything to avoid unnecessarily complicating the API.
-      // folks can use all the same utilities we return in the first place that are bound to the container
+      if (reactRoot === null) {
+        render(wrapUiIfNeeded(rerenderUi), {container, baseElement})
+        // Intentionally do not return anything to avoid unnecessarily complicating the API.
+        // folks can use all the same utilities we return in the first place that are bound to the container
+      } else {
+        reactRoot.render(wrapUiIfNeeded(rerenderUi))
+      }
     },
     asFragment: () => {
       /* istanbul ignore if (jsdom limitation) */
@@ -95,7 +120,15 @@ function cleanup() {
 // maybe one day we'll expose this (perhaps even as a utility returned by render).
 // but let's wait until someone asks for it.
 function cleanupAtContainer(container) {
-  ReactDOM.unmountComponentAtNode(container)
+  const reactRoot = mountedRoots.get(container)
+  if (reactRoot === undefined) {
+    ReactDOM.unmountComponentAtNode(container)
+  } else {
+    act(() => {
+      reactRoot.unmount()
+    })
+  }
+
   if (container.parentNode === document.body) {
     document.body.removeChild(container)
   }

tests/setup-env.js

@@ -1 +1,3 @@
 import '@testing-library/jest-dom/extend-expect'
+// eslint-disable-next-line import/no-extraneous-dependencies
+jest.mock('scheduler', () => require('scheduler/unstable_mock'))