Optimise context state for less GC pressure

- Avoid Context Map() being copied on every element
- Optimise restoring previous context values
- Split legacy context into separate state

Author: kitten

src/__tests__/visitor.test.js

@@ -9,9 +9,13 @@ import { createPortal } from 'react-dom'
 
 import {
   Dispatcher,
-  clearCurrentContextMap,
-  getCurrentContextMap
+  setCurrentContextStore,
+  setCurrentContextMap,
+  getCurrentContextMap,
+  getCurrentContextStore,
+  readContextValue
 } from '../internals'
+
 import { visitElement } from '../visitor'
 
 const {
@@ -23,7 +27,9 @@ let prevDispatcher = null
 beforeEach(() => {
   prevDispatcher = ReactCurrentDispatcher.current
   ReactCurrentDispatcher.current = Dispatcher
-  clearCurrentContextMap()
+
+  setCurrentContextMap({})
+  setCurrentContextStore(new Map())
 })
 
 afterEach(() => {
@@ -75,14 +81,14 @@ describe('visitElement', () => {
       child = visitElement(child, [], () => {})[0]
     }
 
-    expect(getCurrentContextMap().get(Context)).toBe('testA')
+    expect(readContextValue(Context)).toBe('testA')
     expect(leaf).toHaveBeenCalledWith('testA')
 
     for (let i = 0, child = makeChild('testB'); i <= 3 && child; i++) {
       child = visitElement(child, [], () => {})[0]
     }
 
-    expect(getCurrentContextMap().get(Context)).toBe('testB')
+    expect(readContextValue(Context)).toBe('testB')
     expect(leaf).toHaveBeenCalledWith('testB')
   })
 
@@ -107,6 +113,7 @@ describe('visitElement', () => {
 
     expect(queue[0]).toMatchObject({
       contextMap: getCurrentContextMap(),
+      contextStore: getCurrentContextStore(),
       thenable: expect.any(Promise),
       kind: 'frame.lazy',
       type: Test,

src/index.js

@@ -12,9 +12,8 @@ import {
 } from './render'
 
 import {
-  clearCurrentContextMap,
+  setCurrentContextStore,
   setCurrentContextMap,
-  getCurrentContextMap,
   Dispatcher
 } from './internals'
 
@@ -51,12 +50,14 @@ const flushFrames = (queue: Frame[], visitor: Visitor): Promise<void> => {
   })
 }
 
-const defaultVisitor = () => {};
+const defaultVisitor = () => {}
 
 const renderPrepass = (element: Node, visitor?: Visitor): Promise<void> => {
   const queue: Frame[] = []
   let fn = visitor !== undefined ? visitor : defaultVisitor
-  clearCurrentContextMap()
+
+  setCurrentContextMap({})
+  setCurrentContextStore(new Map())
 
   try {
     prevDispatcher = ReactCurrentDispatcher.current

src/internals/__tests__/state.test.js renamed to src/internals/__tests__/context.test.js

@@ -1,45 +1,44 @@
-import { setCurrentContextMap, readContextMap, maskContext } from '../state'
+import {
+  setCurrentContextStore,
+  setCurrentContextMap,
+  readContextValue,
+  maskContext
+} from '../context'
 
-describe('readContextMap', () => {
+describe('readContextValue', () => {
   it('returns values in a Map by key', () => {
     const map = new Map()
     const ctx = {}
-    setCurrentContextMap(map)
-    map.set('key', 'value')
+    setCurrentContextStore(map)
     map.set(ctx, 'value')
-    expect(readContextMap('key')).toBe('value')
-    expect(readContextMap(ctx)).toBe('value')
+    expect(readContextValue(ctx)).toBe('value')
   })
 
   it('returns default values when keys are unknown', () => {
     setCurrentContextMap(new Map())
     const ctx = { _currentValue: 'default' }
-    expect(readContextMap('key')).toBe(undefined)
-    expect(readContextMap(ctx)).toBe('default')
+    expect(readContextValue(ctx)).toBe('default')
   })
 })
 
 describe('maskContext', () => {
   it('supports contextType', () => {
     const map = new Map()
     const ctx = {}
-    setCurrentContextMap(map)
+    setCurrentContextStore(map)
     map.set(ctx, 'value')
     expect(maskContext({ contextType: ctx })).toBe('value')
   })
 
   it('supports no context', () => {
     const map = new Map()
-    setCurrentContextMap(map)
+    setCurrentContextStore(map)
     expect(maskContext({})).toEqual({})
   })
 
   it('supports contextTypes', () => {
-    const map = new Map()
+    const map = { a: 'a', b: 'b', c: 'c' }
     setCurrentContextMap(map)
-    map.set('a', 'a')
-    map.set('b', 'b')
-    map.set('c', 'c')
     expect(maskContext({ contextTypes: { a: null, b: null } })).toEqual({
       a: 'a',
       b: 'b'

src/internals/context.js

@@ -0,0 +1,89 @@
+// @flow
+
+import type {
+  AbstractContext,
+  UserElement,
+  ContextMap,
+  ContextStore
+} from '../types'
+
+type ContextEntry = [AbstractContext, mixed]
+
+let currentContextStore: ContextStore = new Map()
+let currentContextMap: ContextMap = {}
+
+let prevContextMap: void | ContextMap = undefined
+let prevContextEntry: void | ContextEntry = undefined
+
+export const getCurrentContextMap = (): ContextMap => currentContextMap
+export const getCurrentContextStore = (): ContextStore =>
+  new Map(currentContextStore)
+
+export const flushPrevContextMap = (): void | ContextMap => {
+  const prev = prevContextMap
+  prevContextMap = undefined
+  return prev
+}
+
+export const flushPrevContextStore = (): void | ContextEntry => {
+  const prev = prevContextEntry
+  prevContextEntry = undefined
+  return prev
+}
+
+export const restoreContextMap = (prev: ContextMap) => {
+  currentContextMap = prev
+}
+
+export const restoreContextStore = (prev: ContextEntry) => {
+  currentContextStore.set(prev[0], prev[1])
+}
+
+export const setCurrentContextMap = (map: ContextMap) => {
+  prevContextMap = undefined
+  currentContextMap = map
+}
+
+export const setCurrentContextStore = (store: ContextStore) => {
+  prevContextEntry = undefined
+  currentContextStore = store
+}
+
+export const assignContextMap = (map: ContextMap) => {
+  prevContextMap = currentContextMap
+  currentContextMap = Object.assign({}, currentContextMap, map)
+}
+
+export const setContextValue = (context: AbstractContext, value: mixed) => {
+  prevContextEntry = [context, currentContextStore.get(context)]
+  currentContextStore.set(context, value)
+}
+
+export const readContextValue = (context: AbstractContext) => {
+  const value = currentContextStore.get(context)
+  if (value !== undefined) {
+    return value
+  }
+
+  // Return default if context has no value yet
+  return context._currentValue
+}
+
+const emptyContext = {}
+
+export const maskContext = (type: $PropertyType<UserElement, 'type'>) => {
+  const { contextType, contextTypes } = type
+
+  if (contextType) {
+    return readContextValue(contextType)
+  } else if (!contextTypes) {
+    return emptyContext
+  }
+
+  const maskedContext = {}
+  for (const name in contextTypes) {
+    maskedContext[name] = currentContextMap[name]
+  }
+
+  return maskedContext
+}

src/internals/dispatcher.js

@@ -2,7 +2,7 @@
 // Source: https://github.com/facebook/react/blob/c21c41e/packages/react-dom/src/server/ReactPartialRendererHooks.js
 
 import is from 'object-is'
-import { readContextMap } from './state'
+import { readContextValue } from './context'
 
 import type {
   AbstractContext,
@@ -143,12 +143,12 @@ function readContext(context: AbstractContext, _: void | number | boolean) {
   // NOTE: The warning that is used in ReactPartialRendererHooks is obsolete
   // in a prepass, since it'll be caught by a subsequent renderer anyway
   // https://github.com/facebook/react/blob/c21c41e/packages/react-dom/src/server/ReactPartialRendererHooks.js#L215-L223
-  return readContextMap(context)
+  return readContextValue(context)
 }
 
 function useContext(context: AbstractContext, _: void | number | boolean) {
   getCurrentIdentity()
-  return readContextMap(context)
+  return readContextValue(context)
 }
 
 function basicStateReducer<S>(state: S, action: BasicStateAction<S>): S {

src/internals/index.js

@@ -1,4 +1,4 @@
 // @flow
 
+export * from './context'
 export * from './dispatcher'
-export * from './state'

src/internals/state.js

@@ -14,10 +14,12 @@ import type {
 
 import {
   maskContext,
-  forkContextMap,
+  assignContextMap,
   setCurrentIdentity,
   setCurrentContextMap,
-  getCurrentContextMap
+  getCurrentContextMap,
+  setCurrentContextStore,
+  getCurrentContextStore
 } from '../internals'
 
 const createUpdater = () => {
@@ -95,6 +97,7 @@ const createInstance = (type: any, props: DefaultProps) => {
 
 const makeFrame = (type: any, instance: any, thenable: Promise<any>) => ({
   contextMap: getCurrentContextMap(),
+  contextStore: getCurrentContextStore(),
   thenable,
   kind: 'frame.class',
   instance,
@@ -122,11 +125,8 @@ const render = (type: any, instance: any, queue: Frame[]) => {
     typeof instance.getChildContext === 'function'
   ) {
     const childContext = instance.getChildContext()
-    if (childContext) {
-      const contextMap = forkContextMap()
-      for (const name in childContext) {
-        contextMap.set(name, childContext[name])
-      }
+    if (childContext !== null && typeof childContext === 'object') {
+      assignContextMap(childContext)
     }
   }
 
@@ -171,5 +171,6 @@ export const mount = (
 export const update = (queue: Frame[], frame: ClassFrame) => {
   setCurrentIdentity(null)
   setCurrentContextMap(frame.contextMap)
+  setCurrentContextStore(frame.contextStore)
   return render(frame.type, frame.instance, queue)
 }

src/render/classComponent.js

@@ -18,8 +18,10 @@ import {
   maskContext,
   makeIdentity,
   setCurrentIdentity,
-  setCurrentContextMap,
   getCurrentIdentity,
+  setCurrentContextStore,
+  getCurrentContextStore,
+  setCurrentContextMap,
   getCurrentContextMap,
   renderWithHooks,
   setFirstHook,
@@ -32,6 +34,7 @@ const makeFrame = (
   thenable: Promise<any>
 ) => ({
   contextMap: getCurrentContextMap(),
+  contextStore: getCurrentContextStore(),
   id: getCurrentIdentity(),
   hook: getFirstHook(),
   kind: 'frame.hooks',
@@ -90,5 +93,6 @@ export const update = (queue: Frame[], frame: HooksFrame) => {
   setFirstHook(frame.hook)
   setCurrentIdentity(frame.id)
   setCurrentContextMap(frame.contextMap)
+  setCurrentContextStore(frame.contextStore)
   return render(frame.type, frame.props, queue)
 }