fix early returns in getInitialState

Author: Keyan Zhang

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

@@ -106,3 +106,25 @@ var MyComponent4 = React.createClass({
     this.setState({heyoo: 24});
   },
 });
+
+var Loader = React.createClass({
+  getInitialState() {
+    if (this.props.stuff) {
+      return {x: 1};
+    } else if (this.props.thing) {
+      return {x: 2};
+    }
+    switch (this.props.wow) {
+      case 1:
+        return this.props.lol ?
+          {x: 3} :
+          this.whatever(this.props);
+    }
+
+    return this.lol();
+  },
+
+  render() {
+    return null;
+  },
+});

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

@@ -113,3 +113,30 @@ class MyComponent4 extends React.Component {
     this.setState({heyoo: 24});
   };
 }
+
+class Loader extends React.Component {
+  constructor(props, context) {
+    super(props, context);
+    if (props.stuff) {
+      this.state = {x: 1};
+      return;
+    } else if (props.thing) {
+      this.state = {x: 2};
+      return;
+    }
+    switch (props.wow) {
+      case 1:
+        this.state = props.lol ?
+          {x: 3} :
+          this.whatever(props);
+
+        return;
+    }
+
+    this.state = this.lol();
+  }
+
+  render() {
+    return null;
+  }
+}

transforms/class.js

@@ -258,24 +258,34 @@ module.exports = (file, api, options) => {
       })
       .forEach(path => j(path).replaceWith(j.identifier('props')));
 
-  const inlineGetInitialState = getInitialState =>
-    getInitialState.value.body.body.map(statement => {
-      if (statement.type === 'ReturnStatement') {
-        return j.expressionStatement(
+  const inlineGetInitialState = getInitialState => {
+    const functionExpressionAST = j(getInitialState.value);
+
+    return functionExpressionAST
+      .find(j.ReturnStatement)
+      .forEach(path => {
+        j(path).replaceWith(j.expressionStatement(
           j.assignmentExpression(
             '=',
             j.memberExpression(
               j.thisExpression(),
               j.identifier('state'),
               false
             ),
-            statement.argument
+            path.value.argument
           )
-        );
-      }
-
-      return statement;
-    });
+        ));
+
+        if ( // FIXME is there a better way to check this?
+          j(path).closest(j.IfStatement).size() ||
+          j(path).closest(j.SwitchStatement).size() ||
+          j(path).closest(j.WhileStatement).size() ||
+          j(path).closest(j.ForStatement).size()
+        ) {
+          j(path).insertAfter(j.returnStatement(null));
+        }
+      }).getAST()[0].value.body.body;
+  };
 
   const pickReturnValueOrCreateIIFE = value => {
     if (hasSingleReturnStatementWithObject(value)) {