Queue mutate functions before provider mounted

__tests__/index-flowtest.js

@@ -3,21 +3,21 @@
  * @format
  */
 
-import React from 'react';
-import type { Store } from '../src/index.js';
-import createStore from '../src/index.js';
+import React from "react";
+import type { Store } from "../src/index.js";
+import createStore from "../src/index.js";
 
 type User = { name: string, age: number };
 type State = {
   users: Array<User>,
-  version: number,
+  version: number
 };
 
 const store: Store<State> = createStore(
   ({
     users: [],
-    version: 100,
-  }: State),
+    version: 100
+  }: State)
 );
 
 const { Provider, Consumer, mutate } = store;
@@ -79,15 +79,15 @@ function testIncorrectConsumer() {
 
 function testMutate() {
   mutate((draft, state) => {
-    draft.users = [{ name: 'shengmin', age: 20 }];
+    draft.users = [{ name: "shengmin", age: 20 }];
     draft.version = state.version + 1;
   });
 }
 
 function testIncorrectMutate() {
   mutate((draft, state) => {
     // $FlowExpectedError
-    draft.users = '';
+    draft.users = "";
     // $FlowExpectedError
     draft.version = state.versions + 1;
   });

__tests__/index.spec.js

@@ -443,4 +443,42 @@ describe("copy-on-write-store", () => {
     render(<App />);
     expect(log).toEqual(["foo"]);
   });
+
+  it("supports mutate calls before provider mounts", () => {
+    let log = [];
+    const { Provider, Consumer } = createState({ foo: "" });
+
+    let waitForMount = new Promise(resolve => {
+      class MutatingComponent extends React.Component {
+        componentDidMount() {
+          mutate(draft => {
+            draft.foo = "bar";
+            resolve();
+          });
+        }
+
+        render() {
+          return null;
+        }
+      }
+
+      const App = () => (
+        <Provider>
+          <Consumer
+            select={[state => state.foo]}
+            render={foo => {
+              log.push(foo);
+              return null;
+            }}
+          />
+          <MutatingComponent />
+        </Provider>
+      );
+      render(<App />);
+    });
+
+    waitForMount.then(() => {
+      expect(log).toEqual(["bar"]);
+    });
+  });
 });

src/index.js

@@ -15,7 +15,7 @@ import React, { Component } from "react";
 import produce from "immer";
 import invariant from "invariant";
 import shallowEqual from "fbjs/lib/shallowEqual";
-import createContext from 'create-react-context'
+import createContext from "create-react-context";
 
 // The default selector is the identity function
 function identityFn(n) {
@@ -24,16 +24,16 @@ function identityFn(n) {
 
 export default function createCopyOnWriteState(baseState) {
   let updateState = null;
+  let mutateQueue = [];
   const State = createContext(baseState);
   // Wraps immer's produce. Only notifies the Provider
   // if the returned draft has been changed.
   function mutate(fn) {
     invariant(
       updateState !== null,
       `mutate(...): you cannot call mutate when no CopyOnWriteStoreProvider ` +
-        `instance is mounted. Make sure to wrap your consumer components with ` +
-        `the returned Provider, and/or delay your mutate calls until the component ` +
-        `tree is mounted.`
+        `instance is constructed. Make sure to wrap your consumer components with ` +
+        `the returned Provider`
     );
     updateState(fn);
   }
@@ -51,15 +51,33 @@ export default function createCopyOnWriteState(baseState) {
   }
 
   class CopyOnWriteStoreProvider extends React.Component {
+    constructor(props) {
+      super(props);
+
+      // If provider doesn't mounted yet, enqueue requests
+      updateState = mutateQueue.unshift;
+      this.mounted = false;
+    }
+
     state = this.props.initialState || baseState;
 
     componentDidMount() {
       invariant(
-        updateState === null,
+        this.mounted === false,
         `CopyOnWriteStoreProvider(...): There can only be a single ` +
           `instance of a provider rendered at any given time.`
       );
+
+      this.mounted = true;
+
       updateState = this.updateState;
+
+      // dequeue and call requests that pushed the queue before
+      // provider mounted
+      while (mutateQueue.length > 0) {
+        const fn = mutateQueue.pop();
+        updateState(fn);
+      }
     }
 
     componentWillUnmount() {