[Refactor] Stop rerendering directives, inject imports instead

Summary:
This commit changes how the sorted imports are combined with the
original source.

Prior to this commit, all ImportDeclaration nodes and their leading
comments, plus any InterpreterDirective and Directive nodes, were
extracted from the original code and re-rendered using babel. The
rendered nodes were then concatenated with the original source with
those nodes removed to produce the updated source.

This approach safely protected against functional changes, but
removed newlines around comments near the beginning of the file when
the first node of the original source was an ImportDeclaration, as
babel does not preserve whitespace when rendering content.

If a user has configured a plugin that attempts to manage comments
and/or whitespace near the top of the file, such as auto-inserting a
license header (as I am trying to do), this results in conflicts /
formatting churn.

This commit does not directly resolve this incompatibility, however
it better prepares the codebase for a plugin option to be added that
can resolve the issue.

Test Plan:
`yarn install && yarn run test --all`

Note that one snapshot was changed by this commit where a newline was
changed, acting as an effective example of how the original approach
could affect whitespace in the re-rendered portion of the file.

Author: ckwalsh

src/preprocessors/preprocessor.ts

@@ -1,5 +1,5 @@
 import { ParserOptions, parse as babelParser } from '@babel/parser';
-import { Directive, ImportDeclaration } from '@babel/types';
+import { ImportDeclaration } from '@babel/types';
 
 import { PrettierOptions } from '../types';
 import { extractASTNodes } from '../utils/extract-ast-nodes.js';
@@ -27,12 +27,11 @@ export function preprocessor(code: string, options: PrettierOptions) {
     };
 
     const ast = babelParser(code, parserOptions);
-    const interpreter = ast.program.interpreter;
 
     const {
         importNodes,
-        directives,
-    }: { importNodes: ImportDeclaration[]; directives: Directive[] } =
+        injectIdx,
+    }: { importNodes: ImportDeclaration[]; injectIdx: number } =
         extractASTNodes(ast);
 
     // short-circuit if there are no import declaration
@@ -49,7 +48,7 @@ export function preprocessor(code: string, options: PrettierOptions) {
         importOrderSideEffects,
     });
 
-    return getCodeFromAst(allImports, directives, code, interpreter, {
+    return getCodeFromAst(allImports, code, injectIdx, {
         importOrderImportAttributesKeyword,
     });
 }

src/utils/__tests__/assemble-updated-code.spec.ts

@@ -0,0 +1,76 @@
+import { format } from 'prettier';
+
+import { assembleUpdatedCode } from '../assemble-updated-code';
+import { getAllCommentsFromNodes } from '../get-all-comments-from-nodes';
+import { getImportNodes } from '../get-import-nodes';
+import { getSortedNodes } from '../get-sorted-nodes';
+
+const code = `"use strict";
+// first comment
+// second comment
+import z from 'z';
+import c from 'c';
+import g from 'g';
+import t from 't';
+import k from 'k';
+// import a from 'a';
+  // import a from 'a';
+import a from 'a';
+`;
+
+test('it should remove nodes from the original code', async () => {
+    const importNodes = getImportNodes(code);
+    const sortedNodes = getSortedNodes(importNodes, {
+        importOrder: [],
+        importOrderCaseInsensitive: false,
+        importOrderSeparation: false,
+        importOrderGroupNamespaceSpecifiers: false,
+        importOrderSortSpecifiers: false,
+        importOrderSideEffects: true,
+    });
+    const allCommentsFromImports = getAllCommentsFromNodes(sortedNodes);
+
+    const commentAndImportsToRemoveFromCode = [
+        ...sortedNodes,
+        ...allCommentsFromImports,
+    ];
+    const codeWithoutImportDeclarations = assembleUpdatedCode(
+        code,
+        commentAndImportsToRemoveFromCode,
+    );
+    const result = await format(codeWithoutImportDeclarations, {
+        parser: 'babel',
+    });
+    expect(result).toEqual(`"use strict";
+`);
+});
+
+test('it should inject the generated code at the correct location', async () => {
+    const importNodes = getImportNodes(code);
+    const sortedNodes = getSortedNodes(importNodes, {
+        importOrder: [],
+        importOrderCaseInsensitive: false,
+        importOrderSeparation: false,
+        importOrderGroupNamespaceSpecifiers: false,
+        importOrderSortSpecifiers: false,
+        importOrderSideEffects: true,
+    });
+    const allCommentsFromImports = getAllCommentsFromNodes(sortedNodes);
+
+    const commentAndImportsToRemoveFromCode = [
+        ...sortedNodes,
+        ...allCommentsFromImports,
+    ];
+    const codeWithoutImportDeclarations = assembleUpdatedCode(
+        code,
+        commentAndImportsToRemoveFromCode,
+        `import generated from "generated";`,
+        '"use strict";'.length,
+    );
+    const result = await format(codeWithoutImportDeclarations, {
+        parser: 'babel',
+    });
+    expect(result).toEqual(`"use strict";
+import generated from "generated";
+`);
+});

src/utils/__tests__/get-code-from-ast.spec.ts

@@ -1,10 +1,7 @@
-import { ParserOptions, parse as babelParser } from '@babel/parser';
 import { format } from 'prettier';
 import { expect, test } from 'vitest';
 
-import { extractASTNodes } from '../extract-ast-nodes';
 import { getCodeFromAst } from '../get-code-from-ast';
-import { getExperimentalParserPlugins } from '../get-experimental-parser-plugins';
 import { getImportNodes } from '../get-import-nodes';
 import { getSortedNodes } from '../get-sorted-nodes';
 
@@ -28,7 +25,7 @@ import a from 'a';
         importOrderSortByLength: null,
         importOrderSideEffects: true,
     });
-    const formatted = getCodeFromAst(sortedNodes, [], code, null);
+    const formatted = getCodeFromAst(sortedNodes, code);
     expect(await format(formatted, { parser: 'babel' })).toEqual(
         `// first comment
 // second comment
@@ -41,29 +38,3 @@ import z from "z";
 `,
     );
 });
-
-test('it renders directives correctly', async () => {
-    const code = `
-    "use client";
-// first comment
-import b from 'b';
-import a from 'a';`;
-
-    const parserOptions: ParserOptions = {
-        sourceType: 'module',
-        plugins: getExperimentalParserPlugins([]),
-    };
-    const ast = babelParser(code, parserOptions);
-    if (!ast) throw new Error('ast is null');
-    const { directives, importNodes } = extractASTNodes(ast as any);
-
-    const formatted = getCodeFromAst(importNodes, directives, code, null);
-    expect(await format(formatted, { parser: 'babel' })).toEqual(
-        `"use client";
-
-// first comment
-import b from "b";
-import a from "a";
-`,
-    );
-});

src/utils/__tests__/remove-nodes-from-original-code.spec.ts

@@ -0,0 +1,66 @@
+import { Comment, Node } from '@babel/types';
+
+type NodeOrComment = Node | Comment;
+type BoundedNodeOrComment = NodeOrComment & { start: number; end: number };
+
+interface InjectedCode {
+    type: 'InjectedCode';
+    start: number;
+    end: number;
+}
+
+/**
+ * Assembles the updated file, removing imports from the original file and
+ * injecting the sorted imports at the appropriate location.
+ *
+ * @param code the whole file as text
+ * @param nodes to be removed
+ * @param injectedCode the generated import source to be injected
+ * @param injectIdx the index at which to inject the generated source
+ */
+export const assembleUpdatedCode = (
+    code: string,
+    nodes: (Node | Comment)[],
+    injectedCode?: string,
+    injectIdx: number = 0,
+): string => {
+    const ranges: (BoundedNodeOrComment | InjectedCode)[] = nodes.filter(
+        (node): node is BoundedNodeOrComment => {
+            const start = Number(node.start);
+            const end = Number(node.end);
+            return Number.isSafeInteger(start) && Number.isSafeInteger(end);
+        },
+    );
+    if (injectedCode !== undefined) {
+        ranges.push({
+            type: 'InjectedCode',
+            start: injectIdx,
+            end: injectIdx,
+        });
+    }
+    ranges.sort((a, b) => a.start - b.start);
+
+    let result: string = '';
+    let idx = 0;
+
+    for (const { type, start, end } of ranges) {
+        if (start > idx) {
+            result += code.slice(idx, start);
+            idx = start;
+        }
+
+        if (injectedCode !== undefined && type === 'InjectedCode') {
+            result += injectedCode;
+        }
+
+        if (end > idx) {
+            idx = end;
+        }
+    }
+
+    if (idx < code.length) {
+        result += code.slice(idx);
+    }
+
+    return result;
+};

src/utils/assemble-updated-code.ts

@@ -1,26 +1,23 @@
 import { ParseResult } from '@babel/parser';
 import traverseModule, { NodePath } from '@babel/traverse';
-import { Directive, File, ImportDeclaration } from '@babel/types';
+import { Directive, File, ImportDeclaration, Program } from '@babel/types';
 
 const traverse = (traverseModule as any).default || traverseModule;
 
 export function extractASTNodes(ast: ParseResult<File>) {
     const importNodes: ImportDeclaration[] = [];
-    const directives: Directive[] = [];
+    let injectIdx = 0;
     traverse(ast, {
-        Directive(path: NodePath<Directive>) {
-            // Only capture directives if they are at the top scope of the source
-            // and their previous siblings are all directives
-            if (
-                path.parent.type === 'Program' &&
-                path.getAllPrevSiblings().every((s) => {
-                    return s.type === 'Directive';
-                })
-            ) {
-                directives.push(path.node);
-
-                // Trailing comments probably shouldn't be attached to the directive
-                path.node.trailingComments = null;
+        Program(path: NodePath<Program>) {
+            /**
+             * Imports will be injected before the first node of the body and
+             * its comments, skipping InterpreterDirective and Directive nodes.
+             * If the body is empty, default to 0, there will be no imports to
+             * inject anyway.
+             */
+            for (const node of path.node.body) {
+                injectIdx = node.leadingComments?.[0]?.start ?? node.start ?? 0;
+                break;
             }
         },
 
@@ -33,5 +30,5 @@ export function extractASTNodes(ast: ParseResult<File>) {
             }
         },
     });
-    return { importNodes, directives };
+    return { importNodes, injectIdx };
 }

src/utils/extract-ast-nodes.ts

@@ -1,8 +1,9 @@
 import generateModule from '@babel/generator';
-import { Directive, InterpreterDirective, Statement, file } from '@babel/types';
+import { Statement, file } from '@babel/types';
 
 import { newLineCharacters } from '../constants.js';
 import { PrettierOptions } from '../types';
+import { assembleUpdatedCode } from './assemble-updated-code';
 import { getAllCommentsFromNodes } from './get-all-comments-from-nodes.js';
 import { removeNodesFromOriginalCode } from './remove-nodes-from-original-code.js';
 
@@ -15,31 +16,19 @@ const generate = (generateModule as any).default || generateModule;
  */
 export const getCodeFromAst = (
     nodes: Statement[],
-    directives: Directive[],
     originalCode: string,
-    interpreter?: InterpreterDirective | null,
+    injectIdx: number = 0,
     options?: Pick<PrettierOptions, 'importOrderImportAttributesKeyword'>,
 ) => {
     const allCommentsFromImports = getAllCommentsFromNodes(nodes);
 
-    const nodesToRemoveFromCode = [
-        ...directives,
-        ...nodes,
-        ...allCommentsFromImports,
-        ...(interpreter ? [interpreter] : []),
-    ];
-
-    const codeWithoutImportsAndInterpreter = removeNodesFromOriginalCode(
-        originalCode,
-        nodesToRemoveFromCode,
-    );
+    const nodesToRemoveFromCode = [...nodes, ...allCommentsFromImports];
 
     const newAST = file({
         type: 'Program',
         body: nodes,
-        directives,
+        directives: [],
         sourceType: 'module',
-        interpreter: interpreter,
         leadingComments: [],
         innerComments: [],
         trailingComments: [],
@@ -57,10 +46,13 @@ export const getCodeFromAst = (
         importAttributesKeyword: options?.importOrderImportAttributesKeyword,
     });
 
-    return (
+    return assembleUpdatedCode(
+        originalCode,
+        nodesToRemoveFromCode,
         code.replace(
             /"PRETTIER_PLUGIN_SORT_IMPORTS_NEW_LINE";/gi,
             newLineCharacters,
-        ) + codeWithoutImportsAndInterpreter.trim()
+        ),
+        injectIdx,
     );
 };