Merge pull request #4 from codemodsquad/fix/conditional-returns

fix several issues with conditional returns

Author: jedwards1211

src/util/canUnwindAsIs.ts

@@ -0,0 +1,24 @@
+import * as t from '@babel/types'
+import { NodePath } from '@babel/traverse'
+
+export default function canUnwindAsIs(
+  path: NodePath<t.CallExpression>
+): boolean {
+  let parent: NodePath<any> = path.parentPath
+  let child: NodePath<any> = path
+  while (parent && !parent.isFunction()) {
+    if (parent.isBlockStatement()) {
+      const body = (parent as NodePath<t.BlockStatement>).get('body')
+      if (child !== body[body.length - 1]) return false
+    } else if (parent.isLoop()) {
+      return false
+    } else if (parent.isSwitchCase()) {
+      if ((parent.node as t.SwitchCase).test != null) return false
+    } else if (!parent.isStatement() && !parent.isAwaitExpression()) {
+      return false
+    }
+    child = parent
+    parent = child.parentPath
+  }
+  return true
+}

src/util/convertConditionalReturns.ts

@@ -1,6 +1,7 @@
 import * as t from '@babel/types'
 import { NodePath } from '@babel/traverse'
-import restOfBlockStatement from './restOfBlockStatement'
+import replaceWithStatements from './replaceWithStatements'
+import removeRestOfBlockStatement from './removeRestOfBlockStatement'
 
 function isLastStatementInBlock(path: NodePath<any>): boolean {
   const { parentPath } = path
@@ -9,81 +10,74 @@ function isLastStatementInBlock(path: NodePath<any>): boolean {
   return (path as NodePath<any>) === body[body.length - 1]
 }
 
-function isInBranch<T extends t.Statement>(
-  path: NodePath<T>,
-  branch: 'consequent' | 'alternate'
-): boolean {
-  const { parentPath } = path
-  if (parentPath.isIfStatement())
-    return (path as NodePath<any>) === parentPath.get(branch)
-  if (parentPath.isBlockStatement()) {
-    const grandparent = parentPath.parentPath
-    return grandparent.isIfStatement() && parentPath === grandparent.get(branch)
+function hasReturn(path: NodePath<t.Statement>): boolean {
+  if (path.isReturnStatement()) return true
+  if (path.isBlockStatement()) {
+    for (const child of (path as NodePath<t.BlockStatement>).get('body')) {
+      if (child.isReturnStatement()) return true
+    }
   }
   return false
 }
 
-const isInConsequent = <T extends t.Statement>(path: NodePath<T>) =>
-  isInBranch(path, 'consequent')
-const isInAlternate = <T extends t.Statement>(path: NodePath<T>) =>
-  isInBranch(path, 'alternate')
-
-function convertToBlockStatement(
-  blockOrExpression: NodePath<any>
-): NodePath<t.BlockStatement> {
-  if (blockOrExpression.isBlockStatement()) return blockOrExpression
-  return (blockOrExpression.replaceWith(
-    t.blockStatement(
-      blockOrExpression.node == null
-        ? []
-        : [
-            blockOrExpression.isStatement()
-              ? blockOrExpression.node
-              : t.expressionStatement(blockOrExpression.node),
-          ]
+function splitBranches(
+  path: NodePath<t.IfStatement>
+): {
+  returning: NodePath<t.Statement>[]
+  notReturning: (NodePath<t.Statement> | NodePath<null>)[]
+} {
+  const returning: NodePath<t.Statement>[] = []
+  const notReturning: (NodePath<t.Statement> | NodePath<null>)[] = []
+  let p: NodePath<t.IfStatement> | NodePath<null> = path
+  while (p.isIfStatement()) {
+    const consequent = (p as NodePath<t.IfStatement>).get('consequent')
+    const alternate: NodePath<any> = (p as NodePath<t.IfStatement>).get(
+      'alternate'
     )
-  ) as any)[0]
-}
+    ;(hasReturn(consequent) ? returning : notReturning).push(consequent)
+    if (!alternate.isIfStatement()) {
+      ;(hasReturn(alternate) ? returning : notReturning).push(alternate)
+    }
 
-function addRestToConsequent<T extends t.Statement>(path: NodePath<T>): void {
-  const ifStatement = path.findParent(p => p.isIfStatement())
-  if (!ifStatement) throw new Error('failed to find parent IfStatement')
-  const rest = restOfBlockStatement(ifStatement)
-  if (!rest.length) return
-  const consequent = (ifStatement as NodePath<t.IfStatement>).get('consequent')
-  const restNodes = rest.map((path: NodePath<t.Statement>) => path.node)
-  convertToBlockStatement(consequent).pushContainer('body', restNodes)
-  rest.forEach((path: NodePath<t.Statement>) => path.remove())
-}
-
-function addRestToAlternate<T extends t.Statement>(path: NodePath<T>): void {
-  const ifStatement = path.findParent(p => p.isIfStatement())
-  if (!ifStatement) throw new Error('failed to find parent IfStatement')
-  const rest = restOfBlockStatement(ifStatement)
-  if (!rest.length) return
-  let alternate: NodePath<any> = ifStatement
-  while (alternate.isIfStatement()) {
-    alternate = (alternate as NodePath<t.IfStatement>).get('alternate')
+    p = alternate
   }
-  const restNodes = rest.map((path: NodePath<t.Statement>) => path.node)
-  convertToBlockStatement(alternate).pushContainer('body', restNodes)
-  rest.forEach((path: NodePath<t.Statement>) => path.remove())
+  return { returning, notReturning }
 }
 
 export default function convertConditionalReturns(
   parent: NodePath<t.BlockStatement>
 ): boolean {
+  let ifDepth = 0
   let isUnwindable = true
+  const ifStatements: NodePath<t.IfStatement>[] = []
   const returnStatements: NodePath<t.ReturnStatement>[] = []
   parent.traverse(
     {
+      IfStatement: {
+        enter(path: NodePath<t.IfStatement>) {
+          if (path.parentPath.isIfStatement()) return
+          ifDepth++
+          const { returning, notReturning } = splitBranches(path)
+          if (returning.length > 0) {
+            if (notReturning.length === 1) {
+              ifStatements.push(path)
+            } else if (notReturning.length > 1) {
+              isUnwindable = false
+              path.stop()
+              return
+            }
+          }
+        },
+        exit(path: NodePath<t.IfStatement>) {
+          if (path.parentPath.isIfStatement()) return
+          ifDepth--
+        },
+      },
       ReturnStatement(path: NodePath<t.ReturnStatement>) {
         let { parentPath } = path
-        let ifDepth = 0
         let loopDepth = 0
         while (parentPath && parentPath !== parent) {
-          if (parentPath.isIfStatement()) ifDepth++
-          else if (parentPath.isLoop()) loopDepth++
+          if (parentPath.isLoop()) loopDepth++
           if (
             loopDepth > 1 ||
             (!isLastStatementInBlock(parentPath) &&
@@ -104,11 +98,15 @@ export default function convertConditionalReturns(
     parent.state
   )
   if (!isUnwindable) return false
-  let returnStatement
-  while ((returnStatement = returnStatements.pop())) {
-    if (isInConsequent(returnStatement)) addRestToAlternate(returnStatement)
-    else if (isInAlternate(returnStatement))
-      addRestToConsequent(returnStatement)
+  let ifStatement
+  while ((ifStatement = ifStatements.pop())) {
+    const {
+      notReturning: [branch],
+    } = splitBranches(ifStatement)
+    const rest = removeRestOfBlockStatement(ifStatement)
+    if (branch.isBlockStatement())
+      (branch as NodePath<t.BlockStatement>).pushContainer('body', rest)
+    else replaceWithStatements(branch, rest)
   }
   return true
 }

src/util/finalCleanup.ts

@@ -17,9 +17,36 @@ function unwrapPromiseResolves(
   return node as t.Expression
 }
 
+function isEmptyBlock(path: NodePath<any>): boolean {
+  return path.isBlockStatement() && path.node.body.length === 0
+}
+
 export default function finalCleanup(path: NodePath<t.Function>): void {
   path.traverse(
     {
+      IfStatement: {
+        exit(path: NodePath<t.IfStatement>) {
+          const consequent = path.get('consequent')
+          const alternate = path.get('alternate')
+          if (isEmptyBlock(consequent)) {
+            if (alternate.node == null) {
+              path.remove()
+            } else if (isEmptyBlock(alternate)) {
+              path.remove()
+            } else {
+              path.replaceWith(
+                t.ifStatement(
+                  t.unaryExpression('!', path.node.test),
+                  alternate.node
+                )
+              )
+            }
+          } else if (isEmptyBlock(alternate)) {
+            path.node.alternate = null
+            alternate.remove()
+          }
+        },
+      },
       AwaitExpression(path: NodePath<t.AwaitExpression>) {
         const argument = path.get('argument')
         const { parentPath } = path

src/util/findAwaitedExpression.ts

@@ -0,0 +1,35 @@
+import * as t from '@babel/types'
+import { NodePath } from '@babel/traverse'
+import { isPromiseMethodCall } from './predicates'
+
+export default function findNextLinkToUnwind(
+  paths: NodePath | NodePath[],
+  except?: t.CallExpression
+): NodePath<t.CallExpression> | null {
+  if (Array.isArray(paths)) {
+    for (const path of paths) {
+      const result = findNextLinkToUnwind(path, except)
+      if (result) return result
+    }
+    return null
+  }
+  const path = paths
+  if (path.node !== except && isPromiseMethodCall(path.node))
+    return path as NodePath<t.CallExpression>
+  let result: NodePath<t.CallExpression> | null = null
+  path.traverse(
+    {
+      CallExpression(path: NodePath<t.CallExpression>) {
+        if (path.node !== except && isPromiseMethodCall(path.node)) {
+          if (result == null) result = path
+          path.stop()
+        }
+      },
+      Function(path: NodePath<t.Function>) {
+        path.skip()
+      },
+    },
+    path.state
+  )
+  return result
+}

src/util/findNextLinkToUnwind.ts

@@ -3,9 +3,9 @@ import { NodePath } from '@babel/traverse'
 
 import parentStatement from './parentStatement'
 
-export default function restOfBlockStatement(
+export default function removeRestOfBlockStatement(
   path: NodePath<any>
-): NodePath<t.Statement>[] {
+): t.Statement[] {
   const statement = parentStatement(path)
   const blockStatement = statement.parentPath
   if (!blockStatement.isBlockStatement())
@@ -14,5 +14,8 @@ export default function restOfBlockStatement(
   const index = body.indexOf(statement)
   if (index < 0)
     throw new Error('failed to get index of Statement within BlockStatement')
-  return body.slice(index + 1)
+  const rest = body.slice(index + 1)
+  const statements = rest.map(p => p.node)
+  rest.forEach(p => p.remove())
+  return statements
 }

src/util/restOfBlockStatement.ts renamed to src/util/removeRestOfBlockStatement.ts

@@ -3,6 +3,7 @@ import { NodePath } from '@babel/traverse'
 
 import getPreceedingLink from './getPreceedingLink'
 import { isNullish } from './predicates'
+import canUnwindAsIs from './canUnwindAsIs'
 import replaceLink from './replaceLink'
 import renameBoundIdentifiers from './renameBoundIdentifiers'
 import unboundIdentifier from './unboundIdentifier'
@@ -39,7 +40,11 @@ export default function unwindCatch(
   }
   const handlerFunction = handler as NodePath<t.Function>
   const body = handlerFunction.get('body')
-  if (body.isBlockStatement() && !convertConditionalReturns(body)) {
+  if (
+    body.isBlockStatement() &&
+    !canUnwindAsIs(link) &&
+    !convertConditionalReturns(body)
+  ) {
     return getPreceedingLink(link)
   }
 

src/util/unwindCatch.ts

@@ -1,7 +1,6 @@
 import * as t from '@babel/types'
 import { NodePath } from '@babel/traverse'
 
-import findAwaitedExpression from './findAwaitedExpression'
 import getThenHandler from './getThenHandler'
 import getCatchHandler from './getCatchHandler'
 import getFinallyHandler from './getFinallyHandler'
@@ -10,6 +9,7 @@ import { unwindThen } from './unwindThen'
 import unwindFinally from './unwindFinally'
 import parentStatement from './parentStatement'
 import replaceWithImmediatelyInvokedAsyncArrowFunction from './replaceWithImmediatelyInvokedAsyncArrowFunction'
+import findNextLinkToUnwind from './findNextLinkToUnwind'
 
 export default function unwindPromiseChain(
   path: NodePath<t.CallExpression>
@@ -24,10 +24,11 @@ export default function unwindPromiseChain(
 
   const { scope } = parentStatement(path)
 
-  let link: NodePath<t.Expression> | null = path as any
+  let link: NodePath<t.CallExpression> | null = path as any
 
-  while (link && link.isCallExpression()) {
-    const callee = (link as NodePath<t.CallExpression>).get('callee')
+  while (link) {
+    const origNode = link.node
+    const callee = link.get('callee')
     if (!callee.isMemberExpression()) break
 
     const thenHandler = getThenHandler(link)
@@ -42,7 +43,7 @@ export default function unwindPromiseChain(
     } else if (finallyHandler) {
       replacements = unwindFinally(finallyHandler)
     }
-    link = replacements ? findAwaitedExpression(replacements) : null
+    link = replacements ? findNextLinkToUnwind(replacements, origNode) : null
     ;(scope as any).crawl()
   }
 }

src/util/unwindPromiseChain.ts

@@ -4,6 +4,7 @@ import { NodePath } from '@babel/traverse'
 import getPreceedingLink from './getPreceedingLink'
 import { awaited } from './builders'
 import { isNullish } from './predicates'
+import canUnwindAsIs from './canUnwindAsIs'
 import renameBoundIdentifiers from './renameBoundIdentifiers'
 import hasMutableIdentifiers from './hasMutableIdentifiers'
 import prependBodyStatement from './prependBodyStatement'
@@ -24,7 +25,11 @@ export function unwindThen(
     const handlerFunction = handler as NodePath<t.Function>
     const input = handlerFunction.get('params')[0]
     const body = handlerFunction.get('body')
-    if (body.isBlockStatement() && !convertConditionalReturns(body)) {
+    if (
+      body.isBlockStatement() &&
+      !canUnwindAsIs(link) &&
+      !convertConditionalReturns(body)
+    ) {
       return getPreceedingLink(link)
     }