From a6ee5b876a90e87efe9fe2f3605bef9347fa4327 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Tue, 15 Nov 2016 14:41:28 +0000
Subject: [PATCH] Memoize merged child context when possible

We can avoid recreating the merged context object if the context provider work was reused. This is essential to avoid allocations for deep setState() calls.
---
 .../shared/fiber/ReactFiberContext.js         | 37 ++++++++++++-------
 .../fiber/__tests__/ReactIncremental-test.js  |  7 +++-
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/src/renderers/shared/fiber/ReactFiberContext.js b/src/renderers/shared/fiber/ReactFiberContext.js
index fbc478020f..5fd47f5a18 100644
--- a/src/renderers/shared/fiber/ReactFiberContext.js
+++ b/src/renderers/shared/fiber/ReactFiberContext.js
@@ -80,23 +80,32 @@ exports.popContextProvider = function() : void {
 exports.pushContextProvider = function(fiber : Fiber, didPerformWork : boolean) : void {
   const instance = fiber.stateNode;
   const childContextTypes = fiber.type.childContextTypes;
-  const childContext = instance.getChildContext();
 
-  for (let contextKey in childContext) {
-    invariant(
-      contextKey in childContextTypes,
-      '%s.getChildContext(): key "%s" is not defined in childContextTypes.',
-      getComponentName(fiber),
-      contextKey
-    );
-  }
-  if (__DEV__) {
-    const name = getComponentName(fiber);
-    const debugID = 0; // TODO: pass a real ID
-    checkReactTypeSpec(childContextTypes, childContext, 'childContext', name, null, debugID);
+  const memoizedMergedChildContext = instance.__reactInternalMemoizedMergedChildContext;
+  const canReuseMergedChildContext = !didPerformWork && memoizedMergedChildContext != null;
+
+  let mergedContext = null;
+  if (canReuseMergedChildContext) {
+    mergedContext = memoizedMergedChildContext;
+  } else {
+    const childContext = instance.getChildContext();
+    for (let contextKey in childContext) {
+      invariant(
+        contextKey in childContextTypes,
+        '%s.getChildContext(): key "%s" is not defined in childContextTypes.',
+        getComponentName(fiber),
+        contextKey
+      );
+    }
+    if (__DEV__) {
+      const name = getComponentName(fiber);
+      const debugID = 0; // TODO: pass a real ID
+      checkReactTypeSpec(childContextTypes, childContext, 'childContext', name, null, debugID);
+    }
+    mergedContext = {...getUnmaskedContext(), ...childContext};
+    instance.__reactInternalMemoizedMergedChildContext = mergedContext;
   }
 
-  const mergedContext = Object.assign({}, getUnmaskedContext(), childContext);
   index++;
   contextStack[index] = mergedContext;
   didPerformWorkStack[index] = didPerformWork;
diff --git a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
index 14cc145254..6c5e04d010 100644
--- a/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
+++ b/src/renderers/shared/fiber/__tests__/ReactIncremental-test.js
@@ -1822,8 +1822,8 @@ describe('ReactIncremental', () => {
     statefulInst.setState({x: 1});
     ReactNoop.flush();
     expect(ops).toEqual([
-      // TODO: we should be able to reuse the previous child context.
-      'Intl:provide {"locale":"fr"}',
+      // Intl was memoized so we did not need to
+      // either render it or recompute its context.
       'ShowLocaleClass:read {"locale":"fr"}',
       'ShowLocaleFn:read {"locale":"fr"}',
     ]);
@@ -1895,6 +1895,9 @@ describe('ReactIncremental', () => {
     statefulInst.setState({locale: 'gr'});
     ReactNoop.flush();
     expect(ops).toEqual([
+      // Intl is below setState() so it might have been
+      // affected by it. Therefore we re-render and recompute
+      // its child context.
       'Intl:read null',
       'Intl:provide {"locale":"gr"}',
       'ShowLocaleClass:read {"locale":"gr"}',