[Refactor] Stop rerendering directives, inject imports instead

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,79 @@
+import { format } from 'prettier';
+import { expect, test } from 'vitest';
+
+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,
+        importOrderSortByLength: null,
+    });
+    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,
+        importOrderSortByLength: null,
+    });
+    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/assemble-updated-code.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/extract-ast-nodes.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/get-code-from-ast.ts

@@ -1,10 +1,10 @@
 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.js';
 import { getAllCommentsFromNodes } from './get-all-comments-from-nodes.js';
-import { removeNodesFromOriginalCode } from './remove-nodes-from-original-code.js';
 
 const generate = (generateModule as any).default || generateModule;
 
@@ -15,31 +15,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 +45,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,
     );
 };