Preserve directive prologues fixes #5 #7

src/preprocessor.ts

@@ -7,7 +7,7 @@ import { getCodeFromAst } from './utils/get-code-from-ast';
 import { getExperimentalParserPlugins } from './utils/get-experimental-parser-plugins';
 import { getSortedNodes } from './utils/get-sorted-nodes';
 
-export function preprocessor(code: string, options: PrettierOptions) {
+export function preprocessor(code: string, options: PrettierOptions): string {
     const {
         importOrderParserPlugins,
         importOrder,
@@ -24,6 +24,8 @@ export function preprocessor(code: string, options: PrettierOptions) {
     };
 
     const ast = babelParser(code, parserOptions);
+
+    const directives = ast.program.directives;
     const interpreter = ast.program.interpreter;
 
     traverse(ast, {
@@ -38,7 +40,9 @@ export function preprocessor(code: string, options: PrettierOptions) {
     });
 
     // short-circuit if there are no import declaration
-    if (importNodes.length === 0) return code;
+    if (importNodes.length === 0) {
+        return code;
+    }
 
     const allImports = getSortedNodes(importNodes, {
         importOrder,
@@ -48,5 +52,5 @@ export function preprocessor(code: string, options: PrettierOptions) {
         importOrderSortSpecifiers,
     });
 
-    return getCodeFromAst(allImports, code, interpreter);
+    return getCodeFromAst(allImports, code, directives, interpreter);
 }

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

@@ -22,7 +22,7 @@ import a from 'a';
         importOrderGroupNamespaceSpecifiers: false,
         importOrderSortSpecifiers: false,
     });
-    const formatted = getCodeFromAst(sortedNodes, code, null);
+    const formatted = getCodeFromAst(sortedNodes, code, [], undefined);
     expect(format(formatted, { parser: 'babel' })).toEqual(
         `// first comment
 // second comment

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

@@ -1,11 +1,12 @@
+import { parse as babelParser } from '@babel/parser';
 import { format } from 'prettier';
 
 import { getAllCommentsFromNodes } from '../get-all-comments-from-nodes';
 import { getImportNodes } from '../get-import-nodes';
 import { getSortedNodes } from '../get-sorted-nodes';
 import { removeNodesFromOriginalCode } from '../remove-nodes-from-original-code';
 
-const code = `// first comment
+const code = `"some directive";// first comment
 // second comment
 import z from 'z';
 import c from 'c';
@@ -18,6 +19,7 @@ import a from 'a';
 `;
 
 test('it should remove nodes from the original code', () => {
+    const ast = babelParser(code, { sourceType: 'module' });
     const importNodes = getImportNodes(code);
     const sortedNodes = getSortedNodes(importNodes, {
         importOrder: [],
@@ -31,6 +33,7 @@ test('it should remove nodes from the original code', () => {
     const commentAndImportsToRemoveFromCode = [
         ...sortedNodes,
         ...allCommentsFromImports,
+        ...ast.program.directives,
     ];
     const codeWithoutImportDeclarations = removeNodesFromOriginalCode(
         code,

src/utils/get-all-comments-from-nodes.ts

@@ -1,6 +1,6 @@
-import { CommentBlock, CommentLine, Statement } from '@babel/types';
+import { CommentBlock, CommentLine, Directive, Statement } from '@babel/types';
 
-export const getAllCommentsFromNodes = (nodes: Statement[]) =>
+export const getAllCommentsFromNodes = (nodes: (Directive | Statement)[]) =>
     nodes.reduce((acc, node) => {
         if (
             Array.isArray(node.leadingComments) &&

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

@@ -1,26 +1,35 @@
 import generate from '@babel/generator';
-import { InterpreterDirective, Statement, file } from '@babel/types';
+import { Directive, InterpreterDirective, Statement, file } from '@babel/types';
 
 import { newLineCharacters } from '../constants';
 import { getAllCommentsFromNodes } from './get-all-comments-from-nodes';
 import { removeNodesFromOriginalCode } from './remove-nodes-from-original-code';
 
 /**
  * This function generate a code string from the passed nodes.
- * @param nodes all imports
- * @param originalCode
+ * @param nodes All imports, in the sorted order in which they should appear in
+ * the generated code.
+ * @param originalCode The original input code that was passed to this plugin.
+ * @param directives All directive prologues from the original code (e.g.
+ * `"use strict";`).
+ * @param interpreter Optional interpreter directives, if present (e.g.
+ * `#!/bin/node`).
  */
 export const getCodeFromAst = (
     nodes: Statement[],
     originalCode: string,
+    directives: Directive[],
     interpreter?: InterpreterDirective | null,
 ) => {
     const allCommentsFromImports = getAllCommentsFromNodes(nodes);
+    const allCommentsFromDirectives = getAllCommentsFromNodes(directives);
 
     const nodesToRemoveFromCode = [
         ...nodes,
         ...allCommentsFromImports,
+        ...allCommentsFromDirectives,
         ...(interpreter ? [interpreter] : []),
+        ...directives,
     ];
 
     const codeWithoutImportsAndInterpreter = removeNodesFromOriginalCode(
@@ -31,7 +40,7 @@ export const getCodeFromAst = (
     const newAST = file({
         type: 'Program',
         body: nodes,
-        directives: [],
+        directives: directives,
         sourceType: 'module',
         interpreter: interpreter,
         sourceFile: '',

src/utils/get-import-nodes.ts

@@ -2,7 +2,10 @@ import { ParserOptions, parse as babelParser } from '@babel/parser';
 import traverse, { NodePath } from '@babel/traverse';
 import { ImportDeclaration, isTSModuleDeclaration } from '@babel/types';
 
-export const getImportNodes = (code: string, options?: ParserOptions) => {
+export const getImportNodes = (
+    code: string,
+    options?: ParserOptions,
+): ImportDeclaration[] => {
     const importNodes: ImportDeclaration[] = [];
     const ast = babelParser(code, {
         ...options,

src/utils/get-sorted-nodes.ts

@@ -17,7 +17,10 @@ import { getSortedNodesByImportOrder } from './get-sorted-nodes-by-import-order'
  * @param nodes All import nodes that should be sorted.
  * @param options Options to influence the behavior of the sorting algorithm.
  */
-export const getSortedNodes: GetSortedNodes = (nodes, options) => {
+export const getSortedNodes: GetSortedNodes = (
+    nodes,
+    options,
+): ImportOrLine[] => {
     const { importOrderSeparation } = options;
 
     // Split nodes at each boundary between a side-effect node and a

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

@@ -1,20 +1,165 @@
 import {
     CommentBlock,
     CommentLine,
+    Directive,
     ImportDeclaration,
     InterpreterDirective,
     Statement,
 } from '@babel/types';
 
-/** Escapes a string literal to be passed to new RegExp. See: https://stackoverflow.com/a/6969486/480608.
- * @param s the string to escape
+/** A range between a start position (inclusive) and an end position (exclusive). */
+type Range = readonly [start: number, end: number];
+
+/** An optional range between a start position (inclusive) and an end position (exclusive). */
+type OptionalRange = readonly [start: number | null, end: number | null];
+
+/** Compares two range by their start position. */
+function compareRangesByStart(range1: Range, range2: Range): number {
+    return range1[0] < range2[0] ? -1 : range1[0] > range2[0] ? 1 : 0;
+}
+
+/** Type predicate that checks whether a range has a defined start and end. */
+function hasRange(range: OptionalRange): range is Range {
+    return (
+        range[0] !== null &&
+        range[1] !== null &&
+        Number.isSafeInteger(range[0]) &&
+        Number.isSafeInteger(range[1])
+    );
+}
+
+/**
+ * @param range1 One range to check.
+ * @param range2 Another range to check.
+ * @param `true` if both ranges have some overlap. This overlap may consist of
+ * a single point, i.e. `[2, 5)` and `[4, 8)` are considered overlapping.
  */
-const escapeRegExp = (s: string) => s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+function hasOverlap(range1: Range, range2: Range): boolean {
+    return range1[1] > range2[0] && range2[1] > range1[0];
+}
 
 /**
- * Removes imports from original file
- * @param code the whole file as text
- * @param nodes to be removed
+ * Given two ranges that are known to overlap, constructs a new range
+ * representing the single range enclosing both ranges.
+ * @param range1 One range to process.
+ * @param range2 Another range to process.
+ * @returns A single range representing the union of both ranges.
+ */
+function mergeOverlappingRanges(range1: Range, range2: Range): Range {
+    return [Math.min(range1[0], range2[0]), Math.max(range1[1], range2[1])];
+}
+
+/**
+ * Given a list of ordered, possibly overlapping (non-disjoint) ranges,
+ * constructs a new list of ranges that consists entirely of disjoint ranges.
+ * The new list is also ordered.
+ * @param A list of ranges that may be overlapping, but are ordered by their
+ * start position.
+ * @return A list of disjoint ranges that are also ordered by their start
+ * position.
+ */
+function mergeRanges(ranges: Range[]): Range[] {
+    // Start with a result list initialized to the empty list
+    // Iterate over all given ranges. If a range overlaps the last item in
+    // the result list, replace the last item with the merger between that item
+    // and the range. Otherwise, just add the item to the result list.
+    // For comparison, see also
+    // https://www.geeksforgeeks.org/merging-intervals/
+    const merged: Range[] = [];
+    for (const range of ranges) {
+        const currRange = merged[merged.length - 1];
+        if (currRange !== undefined && hasOverlap(currRange, range)) {
+            merged[merged.length - 1] = mergeOverlappingRanges(
+                currRange,
+                range,
+            );
+        } else {
+            merged.push(range);
+        }
+    }
+    return merged;
+}
+
+/**
+ * Takes a list of ordered, disjoint, non-overlapping ranges and a range
+ * `[0, totalLength)` that encloses all those ranges.
+ *
+ * Constructs a new list of ranges representing the negation of the ranges with
+ * respect to the enclosing range `[0, totalLength)`. Put in other words,
+ * subtracts the given ranges from the range `[0, totalLength)`.
+ *
+ * More formally, let `r_1`, `r_2`, ..., `r_n` denote the sets represented by
+ * the given ranges; and let `r` be the set `[0, totalLength)`. Then this
+ * function returns a list of ranges representing the set
+ *
+ * > r \ r_1 \ r_2 \ ... \ r_n
+ *
+ * (where `\` is the set negation operator)
+ * @param ranges  A list of disjoint (non-overlapping) ranges ordered by
+ * their start position.
+ * @param totalLength The end of the enclosing range from which to subtract
+ * the given ranges.
+ * @returns A list of ranges representing the inverse of the given ranges with
+ * respect to the enclosing range.
+ */
+function invertRanges(ranges: Range[], totalLength: number): Range[] {
+    // We'd run into out-of-bounds checks if we performed the rest of the
+    // algorithm with an empty array, and would have to insert additional
+    // checks. So just return immediately to keep the code simpler.
+    if (ranges.length === 0) {
+        return ranges;
+    }
+
+    const resultRanges: Range[] = [];
+    const isValidRange = (start: number, end: number) => end > start;
+
+    // Add the part from the start of the enclosing range to the start of the
+    // first range to exclude.
+    //
+    // |-----------xxxxx-----xxxx-----xxxx-----------|
+    //  ^---------^
+    //  This part
+    const firstRange = ranges[0];
+    if (isValidRange(0, firstRange[0])) {
+        resultRanges.push([0, firstRange[0]]);
+    }
+
+    // Iterate over the parts between the ranges to remove and add those parts.
+    //
+    // |----------xxxxx-----xxxx------xxxx-----------|
+    //                 ^---^    ^----^
+    //                   These parts
+    for (let index = 0; index < ranges.length - 1; index += 1) {
+        const prevRange = ranges[index];
+        const nextRange = ranges[index + 1];
+        const start = prevRange[1];
+        const end = nextRange[0];
+        if (isValidRange(start, end)) {
+            resultRanges.push([start, end]);
+        }
+    }
+
+    // Add the part from the end of the last range to exclude to the end of the
+    // enclosing range.
+    //
+    // |----------xxxxx-----xxxx-----xxxx------------|
+    //                                   ^----------^
+    //                                    This part
+    const lastRange = ranges[ranges.length - 1];
+    if (isValidRange(lastRange[1], totalLength)) {
+        resultRanges.push([lastRange[1], totalLength]);
+    }
+
+    return resultRanges;
+}
+
+/**
+ * Given a piece of code and a list of nodes that appear in that code, removes
+ * all those nodes from the code.
+ * @param code The whole file as text
+ * @param nodes List of nodes to be removed from the code.
+ * @return The given code with all parts of the code removed that correspond to
+ * one of the given nodes.
  */
 export const removeNodesFromOriginalCode = (
     code: string,
@@ -24,23 +169,27 @@ export const removeNodesFromOriginalCode = (
         | CommentLine
         | ImportDeclaration
         | InterpreterDirective
+        | Directive
     )[],
-) => {
-    let text = code;
-    for (const node of nodes) {
-        const start = Number(node.start);
-        const end = Number(node.end);
-        if (Number.isSafeInteger(start) && Number.isSafeInteger(end)) {
-            text = text.replace(
-                // only replace imports at the beginning of the line (ignoring whitespace)
-                // otherwise matching commented imports will be replaced
-                new RegExp(
-                    '^\\s*' + escapeRegExp(code.substring(start, end)),
-                    'm',
-                ),
-                '',
-            );
-        }
-    }
-    return text;
+): string => {
+    // A list of ranges we should remove from the code
+    // Each range [start, end] consists of a start position in the code
+    // (inclusive) and an end position in the code (exclusive)
+    const excludes = nodes
+        .map((node) => [node.start, node.end] as const)
+        .filter(hasRange)
+        .sort(compareRangesByStart);
+
+    // In case there are overlapping ranges, merge all overlapping ranges into
+    // a single range.
+    const mergedExcludes = mergeRanges(excludes);
+
+    // Find the code ranges we want to keep by inverting the excludes with
+    // respect to the entire range [0, code.length]
+    const includes = invertRanges(mergedExcludes, code.length);
+
+    // Extract all code parts we want to keep and join them together again
+    return includes
+        .map((include) => code.substring(include[0], include[1]))
+        .join('');
 };