fix: preserve comments in replacement patterns

Author: jedwards1211

src/backend/parse.ts

@@ -2,6 +2,7 @@ import { Backend } from './Backend'
 import { NodePath, Expression, Statement, Node, Comment } from '../types'
 import ensureArray from '../util/ensureArray'
 import forEachNode from '../util/forEachNode'
+import transferComments from '../util/transferComments'
 
 function parse0(
   backend: Backend,
@@ -14,7 +15,11 @@ function parse0(
       return backend.template.expression(strings, ...quasis)
     if (result.length > 1) return result
     const node = result[0]
-    return node.type === 'ExpressionStatement' ? node.expression : node
+    if (node.type === 'ExpressionStatement') {
+      transferComments(node, node.expression, { leading: true, trailing: true })
+      return node.expression
+    }
+    return node
   } catch (error) {
     // fallthrough
   }

src/compileReplacement/ExportNamedDeclaration.ts

@@ -5,6 +5,7 @@ import {
   ReplaceableMatch,
 } from '.'
 import compileGenericNodeReplacement from './GenericNodeReplacement'
+import transferComments from '../util/transferComments'
 
 export default function compileExportNamedDeclarationReplacement(
   path: NodePath<ExportNamedDeclaration, ExportNamedDeclaration>,
@@ -27,6 +28,7 @@ export default function compileExportNamedDeclarationReplacement(
           )
         }
       }
+      transferComments(path.node, result)
       return result
     },
   }

src/compileReplacement/GenericNodeReplacement.ts

@@ -5,6 +5,7 @@ import compileReplacement, {
   ReplaceableMatch,
 } from './index'
 import indentDebug from '../compileMatcher/indentDebug'
+import transferComments from '../util/transferComments'
 
 export default function compileGenericNodeReplacement(
   path: NodePath,
@@ -53,6 +54,12 @@ export default function compileGenericNodeReplacement(
         if (value !== undefined) result[key] = value
       }
 
+      transferComments(pattern, result, {
+        leading: true,
+        trailing: true,
+        clone: true,
+      })
+
       return result
     },
   }

src/compileReplacement/ImportDeclaration.ts

@@ -5,6 +5,7 @@ import {
   ReplaceableMatch,
 } from '.'
 import compileGenericNodeReplacement from './GenericNodeReplacement'
+import transferComments from '../util/transferComments'
 
 export default function compileImportDeclarationReplacement(
   path: NodePath<ImportDeclaration, ImportDeclaration>,
@@ -27,6 +28,7 @@ export default function compileImportDeclarationReplacement(
           )
         }
       }
+      transferComments(path.node, result)
       return result
     },
   }

src/compileReplacement/Placeholder.ts

@@ -13,6 +13,7 @@ import {
 } from '../compileMatcher/Placeholder'
 import createReplacementConverter, { bulkConvert } from '../convertReplacement'
 import cloneNode from '../util/cloneNode'
+import transferComments from '../util/transferComments'
 export { unescapeIdentifier }
 
 export function compileArrayPlaceholderReplacement(
@@ -25,18 +26,25 @@ export function compileArrayPlaceholderReplacement(
     getArrayPlaceholder(identifier) || getRestPlaceholder(identifier)
   if (arrayPlaceholder && isCapturePlaceholder(arrayPlaceholder)) {
     const convertReplacement = createReplacementConverter(pattern)
+
+    const baseGenerate = (match: ReplaceableMatch): Node | Node[] => {
+      const captures = match.arrayCaptures?.[arrayPlaceholder]
+      if (captures) {
+        return [
+          ...bulkConvert(
+            captures.map((c) => cloneNode(c)),
+            convertReplacement
+          ),
+        ]
+      }
+      return [...bulkConvert(cloneNode(pattern.value), convertReplacement)]
+    }
+
     return {
       generate: (match: ReplaceableMatch): Node | Node[] => {
-        const captures = match.arrayCaptures?.[arrayPlaceholder]
-        if (captures) {
-          return [
-            ...bulkConvert(
-              captures.map((c) => cloneNode(c)),
-              convertReplacement
-            ),
-          ]
-        }
-        return [...bulkConvert(cloneNode(pattern.value), convertReplacement)]
+        const result = baseGenerate(match)
+        transferComments(pattern.node, result)
+        return result
       },
     }
   }
@@ -50,16 +58,22 @@ export default function compilePlaceholderReplacement(
   const placeholder = getPlaceholder(identifier)
   if (placeholder && isCapturePlaceholder(placeholder)) {
     const convertReplacement = createReplacementConverter(pattern)
+    const baseGenerate = (match: ReplaceableMatch): Node | Node[] => {
+      const capture = match.captures?.[placeholder]
+      if (capture) {
+        const clone = cloneNode(capture)
+        const astx = getAstxMatchInfo(capture)
+        if (astx?.subcapture) return convertReplacement(astx.subcapture)
+        return convertReplacement(clone)
+      }
+      return convertReplacement(cloneNode(pattern.value))
+    }
+
     return {
       generate: (match: ReplaceableMatch): Node | Node[] => {
-        const capture = match.captures?.[placeholder]
-        if (capture) {
-          const clone = cloneNode(capture)
-          const astx = getAstxMatchInfo(capture)
-          if (astx?.subcapture) return convertReplacement(astx.subcapture)
-          return convertReplacement(clone)
-        }
-        return convertReplacement(cloneNode(pattern.value))
+        const result = baseGenerate(match)
+        transferComments(pattern.node, result)
+        return result
       },
     }
   }

src/compileReplacement/StringLiteral.ts

@@ -4,6 +4,7 @@ import { CompiledReplacement, ReplaceableMatch } from './'
 import { unescapeIdentifier } from './Placeholder'
 import cloneNode from '../util/cloneNode'
 import * as t from '@babel/types'
+import transferComments from '../util/transferComments'
 
 export default function compileStringLiteralReplacement(
   path: NodePath<StringLiteral>
@@ -14,7 +15,9 @@ export default function compileStringLiteralReplacement(
     return {
       generate: (match: ReplaceableMatch): StringLiteral => {
         const captured = match.stringCaptures?.[placeholder]
-        return captured ? t.stringLiteral(captured) : cloneNode(pattern)
+        const result = captured ? t.stringLiteral(captured) : cloneNode(pattern)
+        transferComments(pattern, result)
+        return result
       },
     }
   }

src/compileReplacement/TemplateLiteral.ts

@@ -4,6 +4,7 @@ import { CompiledReplacement, ReplaceableMatch } from './'
 import { unescapeIdentifier } from './Placeholder'
 import cloneNode from '../util/cloneNode'
 import * as t from '@babel/types'
+import transferComments from '../util/transferComments'
 
 function generateValue(cooked: string): { raw: string; cooked: string } {
   return { raw: cooked.replace(/\\|`|\${/g, '\\$&'), cooked }
@@ -21,12 +22,14 @@ export default function compileTemplateLiteralReplacement(
         return {
           generate: (match: ReplaceableMatch): TemplateLiteral => {
             const captured = match.stringCaptures?.[placeholder]
-            return captured
+            const result = captured
               ? t.templateLiteral(
                   [t.templateElement(generateValue(captured), true)],
                   []
                 )
               : cloneNode(pattern)
+            transferComments(pattern, result)
+            return result
           },
         }
       }

src/replace.ts

@@ -7,6 +7,7 @@ import pipeline from './util/pipeline'
 import lodash from 'lodash'
 const { last } = lodash
 import { SimpleReplacementInterface } from './util/SimpleReplacementCollector'
+import transferComments from './util/transferComments'
 
 export type ReplaceOptions = {
   backend: Backend
@@ -99,8 +100,8 @@ function doReplace(match: Match, replacements: Node[]) {
   )
 
   if (replacements.length) {
-    transferComments(replacedPaths[0], replacements[0], { leading: true })
-    transferComments(last(replacedPaths), last(replacements), {
+    transferComments(replacedPaths[0]?.node, replacements[0], { leading: true })
+    transferComments(last(replacedPaths)?.node, last(replacements), {
       trailing: true,
     })
     replacedPaths[0]?.replace(...replacements)
@@ -110,30 +111,3 @@ function doReplace(match: Match, replacements: Node[]) {
     replacedPaths[i].prune()
   }
 }
-
-function transferComments(
-  from: NodePath | undefined,
-  to: Node | undefined,
-  options: { leading?: boolean; trailing?: boolean }
-) {
-  if (!from || !to) return
-  const node: any = from.node
-  const leading = options.leading
-    ? node.comments?.filter((c: any) => c.leading) || node.leadingComments
-    : undefined
-  if (leading?.length) {
-    const dest = node.comments
-      ? (to as any).comments || ((to as any).comments = [])
-      : (to as any).leadingComments || ((to as any).leadingComments = [])
-    for (const c of leading) dest.push(c)
-  }
-  const trailing = options.trailing
-    ? node.comments?.filter((c: any) => c.trailing) || node.trailingComments
-    : undefined
-  if (trailing?.length) {
-    const dest = node.comments
-      ? (to as any).comments || ((to as any).comments = [])
-      : (to as any).trailingComments || ((to as any).trailingComments = [])
-    for (const c of trailing) dest.push(c)
-  }
-}

src/util/transferComments.ts

@@ -0,0 +1,47 @@
+import { Comment, Node } from '../types'
+
+type Commentable = {
+  comments?: (Comment & { leading?: boolean; trailing?: boolean })[]
+  leadingComments?: Comment[]
+  trailingComments?: Comment[]
+}
+
+export default function transferComments(
+  _from: Node | undefined,
+  _to: Node | Node[] | undefined,
+  options: { leading?: boolean; trailing?: boolean; clone?: boolean } = {
+    leading: true,
+    trailing: true,
+    clone: true,
+  }
+) {
+  if (!_from || !_to) return
+  const from: Commentable = _from as any
+  const leading = options.leading
+    ? from.comments?.filter((c: any) => c.leading) || from.leadingComments
+    : undefined
+  if (leading?.length) {
+    const to: Commentable = (Array.isArray(_to) ? _to[0] : _to) as any
+    const dest = from.comments
+      ? to.comments || (to.comments = [])
+      : to.leadingComments || (to.leadingComments = [])
+    for (const c of leading) dest.push(options.clone ? cloneComment(c) : c)
+  }
+  const trailing = options.trailing
+    ? from.comments?.filter((c: any) => c.trailing) || from.trailingComments
+    : undefined
+  if (trailing?.length) {
+    const to: Commentable = (
+      Array.isArray(_to) ? _to[_to.length - 1] : _to
+    ) as any
+    const dest = from.comments
+      ? to.comments || (to.comments = [])
+      : to.trailingComments || (to.trailingComments = [])
+    for (const c of trailing) dest.push(options.clone ? cloneComment(c) : c)
+  }
+}
+
+function cloneComment(c: Comment): Comment {
+  const { type, value } = c
+  return { type, value }
+}

test/findReplace/binaryExpressionReplacementLeadingComment.ts

@@ -0,0 +1,19 @@
+import { findReplaceTestcase } from '../findReplaceTestcase'
+import dedent from 'dedent-js'
+
+findReplaceTestcase({
+  file: __filename,
+  input: dedent`
+    const a = 1 + 2
+  `,
+  find: dedent`
+    $a + $b 
+  `,
+  replace: dedent`
+    // this is a test 
+    $b + $a
+  `,
+  expectedReplace: dedent`
+    const a = 2 + 1 // this is a test
+  `,
+})

test/findReplace/stringPlaceholderReplacementLeadingComment.ts

@@ -0,0 +1,22 @@
+import { findReplaceTestcase } from '../findReplaceTestcase'
+import dedent from 'dedent-js'
+
+findReplaceTestcase({
+  file: __filename,
+  input: dedent`
+    t.number()
+    t.string('foo')
+  `,
+  find: dedent`
+    t.string('$s')
+  `,
+  replace: dedent`
+    // this is a test 
+    $s
+  `,
+  expectedReplace: dedent`
+    t.number()
+    // this is a test
+    'foo'
+  `,
+})