handle inner function declarations in getInitialState correctly

Author: Keyan Zhang

transforms/__testfixtures__/class-initial-state.input.js

@@ -1,8 +1,6 @@
 import React from 'React';
 
-/*
- * Multiline
- */
+// only needs props
 var MyComponent = React.createClass({
   getInitialState: function() {
     var x = this.props.foo;
@@ -21,7 +19,7 @@ var ComponentWithBothPropsAndContextAccess = React.createClass({
     name: React.PropTypes.string,
   },
 
-  // we actually _don't_ need a constructor here since this will be
+  // we actually don't need a constructor here since this will be
   // initialized after a proper super(props, context) call.
   // in other words, `this` will be ready when it reaches here.
   getInitialState: function() {
@@ -54,7 +52,7 @@ const App = React.createClass({
 const App2 = React.createClass({
   getInitialState() {
     const state = {
-      whatever: this.context.whatever,
+      whatever: this.context.whatever, // needs context
     };
     return state;
   },
@@ -84,7 +82,7 @@ const getContextFromInstance = (x) => x.context; // meh
 
 var MyComponent3 = React.createClass({
   getInitialState: function() {
-    var x = getContextFromInstance(this);
+    var x = getContextFromInstance(this); // `this` is referenced alone
     return {
       heyoo: x,
     };
@@ -95,6 +93,8 @@ var MyComponent3 = React.createClass({
   },
 });
 
+// we are not sure what you'll need from `this`,
+// so it's safe to defer `state`'s initialization
 var MyComponent4 = React.createClass({
   getInitialState: function() {
     return {
@@ -107,6 +107,7 @@ var MyComponent4 = React.createClass({
   },
 });
 
+// intense control flow testing
 var Loader = React.createClass({
   getInitialState() {
     if (this.props.stuff) {
@@ -120,6 +121,17 @@ var Loader = React.createClass({
           {x: 3} :
           this.whatever(this.props);
     }
+    for (let i = 0; i < 100; i++) {
+      if (i === 20) {
+        return {x: i};
+      }
+    }
+
+    try {
+      doSomeThingReallyBad();
+    } catch (e) {
+      return {error: e};
+    }
 
     return this.lol();
   },
@@ -129,6 +141,33 @@ var Loader = React.createClass({
   },
 });
 
+var FunctionDeclarationInGetInitialState = React.createClass({
+  getInitialState() {
+    function func() {
+      var x = 1;
+      return x; // dont change me
+    }
+
+    const foo = () => {
+      return 120; // dont change me
+    };
+
+    var q = function() {
+      return 100; // dont change me
+    };
+
+    return {
+      x: func(),
+      y: foo(),
+      z: q(),
+    };
+  },
+
+  render() {
+    return null;
+  },
+});
+
 var DeferStateInitialization = React.createClass({
   getInitialState() {
     return {x: this.something};
@@ -143,7 +182,7 @@ var DeferStateInitialization = React.createClass({
 
 var helper = () => {};
 
-var PassGetInitialState = React.createClass({
+var PassGetInitialState = React.createClass({ // bail out here
   getInitialState() {
     return this.lol();
   },
@@ -157,7 +196,7 @@ var PassGetInitialState = React.createClass({
   },
 });
 
-var UseGetInitialState = React.createClass({
+var UseGetInitialState = React.createClass({ // bail out here
   getInitialState() {
     return this.lol();
   },
@@ -171,7 +210,7 @@ var UseGetInitialState = React.createClass({
   },
 });
 
-var UseArguments = React.createClass({
+var UseArguments = React.createClass({ // bail out here
   helper() {
     console.log(arguments);
   },
@@ -181,7 +220,7 @@ var UseArguments = React.createClass({
   },
 });
 
-var ShadowingIssue = React.createClass({
+var ShadowingIssue = React.createClass({ // bail out here
   getInitialState() {
     const props = { x: 123 };
     return { x: props.x };
@@ -192,6 +231,7 @@ var ShadowingIssue = React.createClass({
   },
 });
 
+// will remove unnecessary bindings
 var ShadowingButFine = React.createClass({
   getInitialState() {
     const props = this.props;

transforms/__testfixtures__/class-initial-state.output.js

@@ -1,8 +1,6 @@
 import React from 'React';
 
-/*
- * Multiline
- */
+// only needs props
 class MyComponent extends React.Component {
   constructor(props) {
     super(props);
@@ -23,7 +21,7 @@ class ComponentWithBothPropsAndContextAccess extends React.Component {
     name: React.PropTypes.string,
   };
 
-  // we actually _don't_ need a constructor here since this will be
+  // we actually don't need a constructor here since this will be
   // initialized after a proper super(props, context) call.
   // in other words, `this` will be ready when it reaches here.
   state = {
@@ -58,7 +56,7 @@ class App2 extends React.Component {
   constructor(props, context) {
     super(props, context);
     const state = {
-      whatever: this.context.whatever,
+      whatever: this.context.whatever, // needs context
     };
     this.state = state;
   }
@@ -92,7 +90,7 @@ const getContextFromInstance = (x) => x.context; // meh
 class MyComponent3 extends React.Component {
   constructor(props, context) {
     super(props, context);
-    var x = getContextFromInstance(this);
+    var x = getContextFromInstance(this); // `this` is referenced alone
 
     this.state = {
       heyoo: x,
@@ -104,6 +102,8 @@ class MyComponent3 extends React.Component {
   };
 }
 
+// we are not sure what you'll need from `this`,
+// so it's safe to defer `state`'s initialization
 class MyComponent4 extends React.Component {
   foo = (): void => {
     this.setState({heyoo: 24});
@@ -114,6 +114,7 @@ class MyComponent4 extends React.Component {
   };
 }
 
+// intense control flow testing
 class Loader extends React.Component {
   constructor(props, context) {
     super(props, context);
@@ -132,6 +133,19 @@ class Loader extends React.Component {
 
         return;
     }
+    for (let i = 0; i < 100; i++) {
+      if (i === 20) {
+        this.state = {x: i};
+        return;
+      }
+    }
+
+    try {
+      doSomeThingReallyBad();
+    } catch (e) {
+      this.state = {error: e};
+      return;
+    }
 
     this.state = this.lol();
   }
@@ -141,6 +155,34 @@ class Loader extends React.Component {
   }
 }
 
+class FunctionDeclarationInGetInitialState extends React.Component {
+  constructor(props) {
+    super(props);
+    function func() {
+      var x = 1;
+      return x; // dont change me
+    }
+
+    const foo = () => {
+      return 120; // dont change me
+    };
+
+    var q = function() {
+      return 100; // dont change me
+    };
+
+    this.state = {
+      x: func(),
+      y: foo(),
+      z: q(),
+    };
+  }
+
+  render() {
+    return null;
+  }
+}
+
 class DeferStateInitialization extends React.Component {
   something = 42;
   state = {x: this.something};
@@ -152,7 +194,7 @@ class DeferStateInitialization extends React.Component {
 
 var helper = () => {};
 
-var PassGetInitialState = React.createClass({
+var PassGetInitialState = React.createClass({ // bail out here
   getInitialState() {
     return this.lol();
   },
@@ -166,7 +208,7 @@ var PassGetInitialState = React.createClass({
   },
 });
 
-var UseGetInitialState = React.createClass({
+var UseGetInitialState = React.createClass({ // bail out here
   getInitialState() {
     return this.lol();
   },
@@ -180,7 +222,7 @@ var UseGetInitialState = React.createClass({
   },
 });
 
-var UseArguments = React.createClass({
+var UseArguments = React.createClass({ // bail out here
   helper() {
     console.log(arguments);
   },
@@ -190,7 +232,7 @@ var UseArguments = React.createClass({
   },
 });
 
-var ShadowingIssue = React.createClass({
+var ShadowingIssue = React.createClass({ // bail out here
   getInitialState() {
     const props = { x: 123 };
     return { x: props.x };
@@ -201,6 +243,7 @@ var ShadowingIssue = React.createClass({
   },
 });
 
+// will remove unnecessary bindings
 class ShadowingButFine extends React.Component {
   constructor(props, context) {
     super(props, context);

transforms/class.js

@@ -383,23 +383,41 @@ module.exports = (file, api, options) => {
       .forEach(path => j(path).replaceWith(j.identifier('props')));
 
   const inlineGetInitialState = getInitialState => {
-    const functionExpressionAST = j(getInitialState.value);
+    const functionExpressionCollection = j(getInitialState.value);
 
     // at this point if there exists bindings like `const props = ...`, we
     // already know the RHS must be `this.props` (see `isGetInitialStateConstructorSafe`)
     // so it's safe to just remove them
-    functionExpressionAST.find(j.VariableDeclarator, {id: {name: 'props'}})
+    functionExpressionCollection.find(j.VariableDeclarator, {id: {name: 'props'}})
       .forEach(path => j(path).remove());
 
-    functionExpressionAST.find(j.VariableDeclarator, {id: {name: 'context'}})
+    functionExpressionCollection.find(j.VariableDeclarator, {id: {name: 'context'}})
       .forEach(path => j(path).remove());
 
-    return functionExpressionAST
+    return functionExpressionCollection
       .find(j.ReturnStatement)
+      .filter(path => {
+        // filter out inner function declarations here (helper functions, promises, etc.).
+        const mainBodyCollection = j(getInitialState.value.body);
+        return (
+          mainBodyCollection
+            .find(j.ArrowFunctionExpression)
+            .find(j.ReturnStatement, path.value)
+            .size() === 0 &&
+          mainBodyCollection
+            .find(j.FunctionDeclaration)
+            .find(j.ReturnStatement, path.value)
+            .size() === 0 &&
+          mainBodyCollection
+            .find(j.FunctionExpression)
+            .find(j.ReturnStatement, path.value)
+            .size() === 0
+        );
+      })
       .forEach(path => {
         let shouldInsertReturnAfterAssignment = false;
 
-        // if the return statement is not a direct child of the function body
+        // if the return statement is not a direct child of getInitialState's body
         if (getInitialState.value.body.body.indexOf(path.value) === -1) {
           shouldInsertReturnAfterAssignment = true;
         }