Controller as context

src/Controller.js

@@ -16,6 +16,8 @@ export type ControllerState = {
   controller: ?any,
 }
 
+const ControllerContext = React.createContext(null);
+
 class Controller extends React.Component<ControllerProps, ControllerState> {
   controller: any;
 
@@ -42,14 +44,12 @@ class Controller extends React.Component<ControllerProps, ControllerState> {
       return children;
     }
 
-    return React.Children.map(children, (child) => {
-      if (child.type.displayName !== 'Scene') {
-        return child;
-      }
-      const props = {...child.props, controller};
-      return <child.type {...props} />;
-    });
+    return (
+      <ControllerContext.Provider value={controller}>
+        {children}
+      </ControllerContext.Provider>
+    );
   }
 }
 
-export { Controller };
+export { Controller, ControllerContext };

src/Scene.js

@@ -1,5 +1,6 @@
 // @flow
 import { default as React } from 'react';
+import { ControllerContext } from './Controller';
 import ScrollMagic from './lib/scrollmagic';
 import debugAddIndicators from './lib/debug.addIndicators.js';
 
@@ -228,7 +229,7 @@ class SceneBase extends React.PureComponent<SceneBaseProps, SceneBaseState> {
   }
 }
 
-class Scene extends React.PureComponent<SceneProps, {}> {
+class SceneWrapper extends React.PureComponent<SceneProps, {}> {
   static displayName = 'Scene';
 
   render() {
@@ -246,4 +247,12 @@ class Scene extends React.PureComponent<SceneProps, {}> {
   }
 }
 
-export { Scene };
+export const Scene = ({ children, ...props }) => (
+  <ControllerContext.Consumer>
+    {controller => (
+      <SceneWrapper controller={controller} {...props}>
+        {children}
+      </SceneWrapper>
+    )}
+  </ControllerContext.Consumer>
+);