fix: findNextLinkToUnwind bug

fix #16

Author: jedwards1211

src/util/findNextLinkToUnwind.ts

@@ -0,0 +1,31 @@
+import * as t from '@babel/types'
+import { NodePath } from '@babel/traverse'
+
+export default function findNode<N extends t.Node>(
+  path: NodePath<any> | NodePath<any>[],
+  node: N
+): NodePath<any> | null {
+  if (Array.isArray(path)) {
+    for (const p of path) {
+      const result = findNode(p, node)
+      if (result) return result
+    }
+    return null
+  }
+
+  let result: NodePath<any> | null = null
+
+  path.traverse({
+    [node.type](path: NodePath<any>) {
+      if (path.node === node) {
+        result = path
+        path.stop()
+      }
+    },
+    Function(path: NodePath<t.Function>) {
+      path.skip()
+    },
+  })
+
+  return result
+}

src/util/findNode.ts

@@ -11,22 +11,26 @@ import convertBodyToBlockStatement from './convertBodyToBlockStatement'
 import mergeCatchIntoTryFinally from './mergeCatchIntoFinally'
 import { awaited } from './builders'
 import convertConditionalReturns from './convertConditionalReturns'
+import findNode from './findNode'
 
 export default function unwindCatch(
   handler: NodePath<t.Expression>
-): NodePath<any> | NodePath<any>[] {
+): NodePath<any> | null {
   const link = handler.parentPath as NodePath<t.CallExpression>
+  let preceedingLink
   let preceeding
   if (link.node.arguments.length === 2) {
-    preceeding = t.awaitExpression(
-      t.callExpression(link.node.callee, [link.node.arguments[0]])
-    )
+    preceedingLink = t.callExpression(link.node.callee, [
+      link.node.arguments[0],
+    ])
+    preceeding = t.awaitExpression(preceedingLink)
   } else {
-    preceeding = awaited(getPreceedingLink(link).node)
+    preceedingLink = getPreceedingLink(link).node
+    preceeding = awaited(preceedingLink)
   }
 
   if (isNullish(handler.node)) {
-    return replaceLink(link, preceeding)
+    return findNode(replaceLink(link, preceeding), preceedingLink)
   }
 
   if (!handler.isFunction()) {
@@ -46,7 +50,7 @@ export default function unwindCatch(
       !isPromiseMethodCall(getPreceedingLink(link).node)) ||
       (!canUnwindAsIs(link) && !convertConditionalReturns(body)))
   ) {
-    return []
+    return null
   }
 
   handlerFunction.node.async = true
@@ -79,8 +83,9 @@ export default function unwindCatch(
   const finalBody = handlerFunction.get('body') as NodePath<t.BlockStatement>
   ;(finalBody.scope as any).crawl()
   const tryStatement = finalBody.get('body')[0] as NodePath<t.TryStatement>
-  return (
+  return findNode(
     mergeCatchIntoTryFinally(link, tryStatement) ||
-    (replaceLink(link, finalBody) as any)
+      (replaceLink(link, finalBody) as any),
+    preceedingLink
   )
 }

src/util/unwindCatch.ts

@@ -9,15 +9,17 @@ import replaceReturnStatements from './replaceReturnStatements'
 import convertBodyToBlockStatement from './convertBodyToBlockStatement'
 import convertConditionalReturns from './convertConditionalReturns'
 import mergeStatementsIntoTryFinally from './mergeStatementsIntoTryFinally'
+import findNode from './findNode'
 
 export default function unwindFinally(
   handler: NodePath<t.Expression>
-): NodePath<any> | NodePath<any>[] {
+): NodePath<any> | null {
   const link = handler.parentPath as NodePath<t.CallExpression>
-  const preceeding = awaited(getPreceedingLink(link).node)
+  const preceedingLink = getPreceedingLink(link)
+  const preceeding = awaited(preceedingLink.node)
 
   if (isNullish(handler.node)) {
-    return replaceLink(link, preceeding)
+    return findNode(replaceLink(link, preceeding), preceedingLink.node)
   }
 
   if (!handler.isFunction()) {
@@ -30,7 +32,7 @@ export default function unwindFinally(
   handlerFunction.node.async = true
   const body = handlerFunction.get('body')
   if (body.isBlockStatement() && !convertConditionalReturns(body)) {
-    return getPreceedingLink(link)
+    return preceedingLink
   }
   body.replaceWith(
     t.blockStatement([
@@ -50,8 +52,9 @@ export default function unwindFinally(
   const finalBody = handlerFunction.get('body')
   ;(finalBody.scope as any).crawl()
   const tryStatement = finalBody.get('body.0') as NodePath<t.TryStatement>
-  return (
+  return findNode(
     mergeStatementsIntoTryFinally(link, tryStatement) ||
-    (replaceLink(link, finalBody) as any)
+      (replaceLink(link, finalBody) as any),
+    preceedingLink.node
   )
 }

src/util/unwindFinally.ts

@@ -9,7 +9,6 @@ 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>
@@ -27,23 +26,22 @@ export default function unwindPromiseChain(
   let link: NodePath<t.CallExpression> | null = path as any
 
   while (link) {
-    const origNode = link.node
     const callee = link.get('callee')
     if (!callee.isMemberExpression()) break
 
     const thenHandler = getThenHandler(link)
     const catchHandler = getCatchHandler(link)
     const finallyHandler = getFinallyHandler(link)
 
-    let replacements: NodePath | NodePath[] | null = null
     if (catchHandler) {
-      replacements = unwindCatch(catchHandler)
+      link = unwindCatch(catchHandler)
     } else if (thenHandler) {
-      replacements = unwindThen(thenHandler)
+      link = unwindThen(thenHandler)
     } else if (finallyHandler) {
-      replacements = unwindFinally(finallyHandler)
+      link = unwindFinally(finallyHandler)
+    } else {
+      link = null
     }
-    link = replacements ? findNextLinkToUnwind(replacements, origNode) : null
     ;(scope as any).crawl()
   }
 }

src/util/unwindPromiseChain.ts

@@ -10,15 +10,17 @@ import hasMutableIdentifiers from './hasMutableIdentifiers'
 import prependBodyStatement from './prependBodyStatement'
 import replaceLink from './replaceLink'
 import convertConditionalReturns from './convertConditionalReturns'
+import findNode from './findNode'
 
 export function unwindThen(
   handler: NodePath<t.Expression>
-): NodePath<any> | NodePath<any>[] {
+): NodePath<any> | null {
   const link = handler.parentPath as NodePath<t.CallExpression>
-  const preceeding = awaited(getPreceedingLink(link).node)
+  const preceedingLink = getPreceedingLink(link)
+  const preceeding = awaited(preceedingLink.node)
 
   if (isNullish(handler.node)) {
-    return replaceLink(link, preceeding)
+    return findNode(replaceLink(link, preceeding), preceedingLink.node)
   }
 
   if (handler.isFunction()) {
@@ -31,7 +33,7 @@ export function unwindThen(
       !canUnwindAsIs(link) &&
       !convertConditionalReturns(body)
     ) {
-      return getPreceedingLink(link)
+      return preceedingLink
     }
 
     if (input) renameBoundIdentifiers(input, link.scope)
@@ -52,7 +54,13 @@ export function unwindThen(
         prepended.get('declarations.0.id') as any
       )
     }
-    return replaceLink(link, handlerFunction.get('body')) as any
+    return findNode(
+      replaceLink(link, handlerFunction.get('body')) as any,
+      preceedingLink.node
+    )
   }
-  return replaceLink(link, t.callExpression(handler.node, [preceeding])) as any
+  return findNode(
+    replaceLink(link, t.callExpression(handler.node, [preceeding])) as any,
+    preceedingLink.node
+  )
 }

src/util/unwindThen.ts

@@ -108,18 +108,11 @@ it('should support many levels of belongsTo', async function() {
 
   await this.sequelize.sync({ force: true });
 
-  await A.bulkCreate([
-    {},
-    {},
-    {},
-    {},
-    {},
-    {},
-    {},
-    {}
-  ]);
-
-  const [as0, b] = await Promise.all([await A.findAll(), (function(singles) {
+  const [as0, b] = await Promise.all([
+    A.bulkCreate([{}, {}, {}, {}, {}, {}, {}, {}]).then(() => {
+      return A.findAll()
+    }),
+    (function(singles) {
     let promise = Promise.resolve(),
       previousInstance,
       b;