Capture current element in rendering

lib/constants.js

@@ -9,5 +9,6 @@
 module.exports = {
 	INTERRUPT_SUSPENSE: 'suspense',
 	INTERRUPT_PROMISE: 'promise',
-	PLACEHOLDER: '__interrupt__'
+	PLACEHOLDER: '__interrupt__',
+	SHIMMED: Symbol('shimmed')
 };

lib/renderer.js

@@ -6,15 +6,69 @@
 'use strict';
 
 // Modules
-const {renderToNodeStream} = require('react-dom/server');
+const React = require('react'),
+	{renderToNodeStream} = require('react-dom/server');
 
 // Imports
-const {INTERRUPT_SUSPENSE, INTERRUPT_PROMISE, PLACEHOLDER} = require('./constants'),
-	{isPromise, isSuspense} = require('./utils');
+const {INTERRUPT_SUSPENSE, INTERRUPT_PROMISE, PLACEHOLDER, SHIMMED} = require('./constants'),
+	{isObject, isPromise, isSuspense, last} = require('./utils');
+
+/*
+ * Shim `React.isValidElement()` to capture elements resolved in render.
+ * Renderer's `.render()` method calls private `resolve()` function inside React.
+ * `resolve()` iterates through user-defined React components, calling each's
+ * `.render()` method and then resolving the element returned in turn.
+ * This shim keeps track of the current element being rendered
+ * (i.e. the one returned by `resolve()`).
+ * This shim also shims the `.getChildContext()` prototype method of any
+ * components in the tree to capture any legacy context emitted as tree is
+ * iterated through in `resolve()`.
+ */
+let inRender = false, currentElement = null, currentContextStack = [];
+
+const {isValidElement} = React;
+React.isValidElement = function(object) {
+	const isElement = isValidElement(object);
+
+	if (inRender && isElement) {
+		currentElement = object;
+		shimGetChildContext(object);
+	}
+
+	return isElement;
+};
+
+function shimGetChildContext(element) {
+	const {type} = element;
+	if (typeof type !== 'function') return;
+
+	const {prototype} = type;
+	if (!isObject(prototype)) return;
+
+	const {getChildContext} = prototype;
+	if (typeof getChildContext !== 'function' || getChildContext[SHIMMED]) return;
+
+	const getChildContextShim = function() {
+		const context = getChildContext.call(this);
+		if (inRender) currentContextStack.push(context);
+		return context;
+	};
+	getChildContextShim[SHIMMED] = true;
+
+	prototype.getChildContext = getChildContextShim;
+}
 
 // Capture ReactDOMServerRenderer class from React
 const ReactDOMServerRenderer = renderToNodeStream('').partialRenderer.constructor;
 
+// Identify if running in dev mode
+const isDev = !!(new ReactDOMServerRenderer('')).stack[0].debugElementStack;
+
+// Function to identify Suspense errors
+const isSuspenseError = isDev ?
+	err => err.message === 'ReactDOMServer does not yet support Suspense.' :
+	err => /^Minified React error #294;/.test(err.message);
+
 // Exports
 class PartialRenderer extends ReactDOMServerRenderer {
 	constructor(children, makeStaticMarkup, stackState) {
@@ -48,7 +102,7 @@ class PartialRenderer extends ReactDOMServerRenderer {
 				context: stackState.context,
 				footer: ''
 			};
-			if (topFrame.debugElementStack) frame.debugElementStack = [];
+			if (isDev) frame.debugElementStack = [];
 
 			stack[1] = frame;
 			topFrame.children = [];
@@ -64,15 +118,31 @@ class PartialRenderer extends ReactDOMServerRenderer {
 	}
 
 	render(element, context, parentNamespace) {
-		if (isSuspense(element)) {
-			return this._interrupt(INTERRUPT_SUSPENSE, element, context, parentNamespace);
-		}
+		inRender = true;
 
 		try {
 			return super.render(element, context, parentNamespace);
 		} catch (err) {
+			// Add captured contexts to context
+			if (currentContextStack.length > 0) {
+				currentContextStack.unshift({}, context);
+				context = Object.assign.apply(null, currentContextStack);
+			}
+
+			if (isSuspense(currentElement) && isSuspenseError(err)) {
+				// Remove last entry from debug stack
+				if (isDev) last(this.stack).debugElementStack.pop();
+
+				return this._interrupt(INTERRUPT_SUSPENSE, currentElement, context, parentNamespace);
+			}
+
 			if (!isPromise(err)) throw err;
+
 			return this._interrupt(INTERRUPT_PROMISE, element, context, parentNamespace, {promise: err});
+		} finally {
+			inRender = false;
+			currentElement = null;
+			currentContextStack.length = 0;
 		}
 	}
 

lib/utils.js

@@ -9,14 +9,32 @@
 const React = require('react');
 
 // Exports
-const {$$typeof: SUSPENSE_TYPEOF, type: SUSPENSE_TYPE} = React.createElement(React.Suspense, null);
-
 module.exports = {
-	isPromise(o) {
-		return o != null && typeof o.then === 'function';
-	},
-
-	isSuspense(e) {
-		return e && e.$$typeof === SUSPENSE_TYPEOF && e.type === SUSPENSE_TYPE;
-	}
+	isObject,
+	isPromise,
+	isReactElement,
+	isSuspense,
+	last
 };
+
+function isObject(o) {
+	return o != null && typeof o === 'object';
+}
+
+function isPromise(o) {
+	return isObject(o) && typeof o.then === 'function';
+}
+
+const {$$typeof: ELEMENT_TYPE, type: SUSPENSE_TYPE} = React.createElement(React.Suspense, null);
+
+function isReactElement(e) {
+	return isObject(e) && e.$$typeof === ELEMENT_TYPE;
+}
+
+function isSuspense(e) {
+	return isReactElement(e) && e.type === SUSPENSE_TYPE;
+}
+
+function last(arr) {
+	return arr[arr.length - 1];
+}