atmn 2.0.21 → 2.0.23

package/src/surgery/appendToArray.ts

@@ -5,51 +5,76 @@ import { leadingIndentOfLine, lineStartOf } from "./fixtureEdit";
 export type AppendText = string | ((elementIndent: string) => string);
 
 /** Append one element to an array literal, matching its siblings' indent and
- * comma style; an empty array is reflowed onto its own lines. */
+ * comma style; an empty array is reflowed onto its own lines. With `after`,
+ * the element lands right behind the last sibling the predicate accepts. */
 export const appendElementToArray = ({
 	source,
 	root,
 	array,
 	text,
+	after,
 }: {
 	source: string;
 	root: SgNode;
 	array: SgNode;
 	text: AppendText;
+	after?: (element: SgNode) => boolean;
 }): string => {
-	const elements = array.namedChildren();
+	// A comment is a named child too, but never an element.
+	const elements = array
+		.namedChildren()
+		.filter((child) => child.kind() !== "comment");
 	if (elements.length === 0) {
 		const indent = leadingIndentOfLine(source, array.range().start.index);
 		// The seeded element's first line lands one tab deeper than the array.
 		const resolved = resolveText({ text, elementIndent: `${indent}\t` });
 		return root.commitEdits([
-			{
-				startPos: array.range().start.index,
-				endPos: array.range().end.index,
-				insertedText: `[\n${indent}\t${resolved},\n${indent}]`,
-			},
+			rebuildArrayEdit({
+				array,
+				lines: [...commentLines({ array }), `${resolved},`],
+				indent,
+			}),
 		]);
 	}
 	const last = elements[elements.length - 1];
-	const lastEnd = last.range().end.index;
-	const after = source.slice(lastEnd);
-	const commaAfter = after.indexOf(",");
-	const hasTrailingComma =
-		commaAfter !== -1 && after.slice(0, commaAfter).trim() === "";
-	const insertAt = hasTrailingComma ? lastEnd + commaAfter + 1 : lastEnd;
-	const missingComma = hasTrailingComma ? "" : ",";
+	let anchorIndex = -1;
+	if (after !== undefined) {
+		for (let index = elements.length - 1; index >= 0; index -= 1) {
+			if (after(elements[index])) {
+				anchorIndex = index;
+				break;
+			}
+		}
+	}
+	const anchor = anchorIndex === -1 ? last : elements[anchorIndex];
+	const { insertAt, hasTrailingComma } = insertionAfter({
+		source,
+		array,
+		anchor,
+	});
+	// A missing comma goes straight after the anchor, ahead of any comment.
+	const commaEdits = hasTrailingComma
+		? []
+		: [
+				{
+					startPos: anchor.range().end.index,
+					endPos: anchor.range().end.index,
+					insertedText: ",",
+				},
+			];
 	const spansLines = source
 		.slice(array.range().start.index, last.range().start.index)
 		.includes("\n");
 	if (spansLines) {
-		// Sibling indent: the whitespace opening the last element's line.
-		const indent = leadingIndentOfLine(source, last.range().start.index);
+		// Sibling indent: the whitespace opening the anchor element's line.
+		const indent = leadingIndentOfLine(source, anchor.range().start.index);
 		const resolved = resolveText({ text, elementIndent: indent });
 		return root.commitEdits([
+			...commaEdits,
 			{
 				startPos: insertAt,
 				endPos: insertAt,
-				insertedText: `${missingComma}\n${indent}${resolved},`,
+				insertedText: `\n${indent}${resolved},`,
 			},
 		]);
 	}
@@ -58,26 +83,89 @@ export const appendElementToArray = ({
 	const resolved = resolveText({ text, elementIndent });
 	if (!resolved.includes("\n")) {
 		return root.commitEdits([
+			...commaEdits,
 			{
 				startPos: insertAt,
 				endPos: insertAt,
-				insertedText: `${missingComma} ${resolved},`,
+				insertedText: ` ${resolved},`,
 			},
 		]);
 	}
-	// A one-line array cannot hold a multi-line element inline: reflow it.
-	const lines = [...elements.map((element) => element.text()), resolved].map(
-		(element) => `${elementIndent}${element},`,
+	// A one-line array cannot hold a multi-line element inline: reflow it,
+	// one line per element, with comments kept in place on their own lines.
+	const children = array.namedChildren();
+	const lines = children.map((child) =>
+		child.kind() === "comment" ? child.text() : `${child.text()},`,
 	);
+	const anchorStart = anchor.range().start.index;
+	const position =
+		anchorIndex === -1
+			? lines.length
+			: children.findIndex(
+					(child) => child.range().start.index === anchorStart,
+				) + 1;
+	lines.splice(position, 0, `${resolved},`);
 	return root.commitEdits([
-		{
-			startPos: array.range().start.index,
-			endPos: array.range().end.index,
-			insertedText: `[\n${lines.join("\n")}\n${lineIndent}]`,
-		},
+		rebuildArrayEdit({ array, lines, indent: lineIndent }),
 	]);
 };
 
+const commentLines = ({ array }: { array: SgNode }): string[] =>
+	array
+		.namedChildren()
+		.filter((child) => child.kind() === "comment")
+		.map((child) => child.text());
+
+/** The whole literal on its own lines, each entry one tab deeper than the array. */
+const rebuildArrayEdit = ({
+	array,
+	lines,
+	indent,
+}: {
+	array: SgNode;
+	lines: string[];
+	indent: string;
+}) => ({
+	startPos: array.range().start.index,
+	endPos: array.range().end.index,
+	insertedText: `[\n${lines.map((line) => `${indent}\t${line}`).join("\n")}\n${indent}]`,
+});
+
+/** Where the next element goes: past the anchor's comma and any comment that
+ * shares its line, so `plan() /* note *\/,` and `plan(), // legacy` keep their trivia. */
+const insertionAfter = ({
+	source,
+	array,
+	anchor,
+}: {
+	source: string;
+	array: SgNode;
+	anchor: SgNode;
+}): { insertAt: number; hasTrailingComma: boolean } => {
+	const children = array.children();
+	const start = children.findIndex(
+		(child) => child.range().start.index === anchor.range().start.index,
+	);
+	let insertAt = anchor.range().end.index;
+	let hasTrailingComma = false;
+	for (const child of children.slice(start + 1)) {
+		const kind = child.kind();
+		if (kind === ",") {
+			if (hasTrailingComma) break;
+			hasTrailingComma = true;
+			insertAt = child.range().end.index;
+			continue;
+		}
+		if (kind !== "comment") break;
+		const sameLine = !source
+			.slice(insertAt, child.range().start.index)
+			.includes("\n");
+		if (!sameLine) break;
+		insertAt = child.range().end.index;
+	}
+	return { insertAt, hasTrailingComma };
+};
+
 const resolveText = ({
 	text,
 	elementIndent,

package/src/surgery/appendToBinding.ts

@@ -1,3 +1,4 @@
+import type { SgNode } from "@ast-grep/napi";
 import { Lang, parse } from "@ast-grep/napi";
 import { type AppendText, appendElementToArray } from "./appendToArray";
 import { findArrayBinding } from "./arrayBinding";
@@ -7,13 +8,15 @@ export const appendToBinding = ({
 	source,
 	name,
 	text,
+	after,
 }: {
 	source: string;
 	name: string;
 	text: AppendText;
+	after?: (element: SgNode) => boolean;
 }): string | null => {
 	const root = parse(Lang.TypeScript, source).root();
 	const array = findArrayBinding({ root, name });
 	if (array === null) return null;
-	return appendElementToArray({ source, root, array, text });
+	return appendElementToArray({ source, root, array, text, after });
 };

package/src/surgery/appendToCollection.ts

@@ -6,10 +6,12 @@ export const appendToCollection = ({
 	source,
 	collection,
 	text,
+	after,
 }: {
 	source: string;
 	collection: string;
 	text: AppendText;
+	after?: (element: SgNode) => boolean;
 }): string | null => {
 	const root = parse(Lang.TypeScript, source).root();
 	// A bare `collection: [$$$]` parses as a type annotation, so the anchor must
@@ -20,7 +22,7 @@ export const appendToCollection = ({
 	if (anchor === null) return null;
 	const array = collectionArray({ anchor, collection });
 	if (array === null) return null;
-	return appendElementToArray({ source, root, array, text });
+	return appendElementToArray({ source, root, array, text, after });
 };
 
 const collectionArray = ({