import traverser from '../traverser';
import {matchesAst, matchesLength, extract} from '../utils/matchesAst';
import copyComments from '../utils/copyComments';

export default function(ast) {
  traverser.replace(ast, {
    enter(node) {
      if (isArrowFunction(node)) {
        return arrowReturn(node);
      }
    }
  });
}

function arrowReturn(node) {
  node.body = extractArrowBody(node.body);
  return node;
}

function isArrowFunction(node) {
  return node.type === 'ArrowFunctionExpression';
}

const matchesReturnBlock = matchesAst({
  type: 'BlockStatement',
  body: matchesLength([
    extract('returnStatement', {
      type: 'ReturnStatement',
      argument: extract('returnVal')
    })
  ])
});

function extractArrowBody(block) {
  const {returnStatement, returnVal} = matchesReturnBlock(block) || {};
  if (returnVal) {
    // preserve return statement comments
    copyComments({from: returnStatement, to: returnVal});
    return returnVal;
  }
  else {
    return block;
  }
}

@@ -242,16 +242,27 @@
   if (is
+Shortenable
 ArrowFun
@@ -291,21 +291,23 @@
  return 
-arrow
+shorten
 Return(n
@@ -348,13 +348,15 @@
 ion 
-arrow
+shorten
 Retu
@@ -431,94 +431,8 @@
 %0A%7D%0A%0A
-function isArrowFunction(node) %7B%0A  return node.type === 'ArrowFunctionExpression';%0A%7D%0A%0A
 cons
@@ -635,16 +635,150 @@
 %5D)%0A%7D);%0A%0A
+function isShortenableArrowFunction(node) %7B%0A  return node.type === 'ArrowFunctionExpression' &&%0A    matchesReturnBlock(node.body);%0A%7D%0A%0A
 function
@@ -872,37 +872,10 @@
 ock)
- %7C%7C %7B%7D;%0A  if (returnVal) %7B%0A  
+;%0A
   //
@@ -910,18 +910,16 @@
 omments%0A
-  
   copyCo
@@ -966,18 +966,16 @@
 nVal%7D);%0A
-  
   return
@@ -990,41 +990,6 @@
 al;%0A
-  %7D%0A  else %7B%0A    return block;%0A  %7D%0A
 %7D%0A

Simplify arrow-return logic