Move pure comment annotation to Graph.contextParse (#3981)

* Move pure comment annotation to Graph.contextParse

If a plugin calls `PluginContext.parse` method to parse the AST of the module
rollup ignored the `PURE` annotations. This commit fixes the issue by moving
the annotation functionality to a central location - the `contextParse` method.

Resolves #3979

* Remove unused import

* Add test call-marked-pure-with-plugin-parse-ast

* Improv test call-marked-pure-with-plugin-parse-ast

Use the `onComment` feautre. This expands the coverage.

* Change type CommentDescription -> acorn.Comment

This change eliminates needless mapping between the types.

* Rename annotations in acron.Node

When parsing the JS code we add attributes to acorn.Node for later use:
* annotations: acorn.Commnet - the comment that defined the annotation
* annotatedPure: boolean - declares this node as pure call/new expression

This commit puts all these attributes under one name: `_rollupAnnotations` with
a strong type.

Thus we clarify that these attributes are for internal use only and decrease
the change of ever coliding with internal/future attributes of `acorn.Node`.

* Remove '_rollupAnnotations' key from Node.keys

* Remove dead code

Co-authored-by: Yannay Livneh <you@example.com>

src/Graph.ts

@@ -17,6 +17,7 @@ import { BuildPhase } from './utils/buildPhase';
 import { errImplicitDependantIsNotIncluded, error } from './utils/error';
 import { analyseModuleExecution } from './utils/executionOrder';
 import { PluginDriver } from './utils/PluginDriver';
+import { markPureCallExpressions } from './utils/pureComments';
 import relativeId from './utils/relativeId';
 import { timeEnd, timeStart } from './utils/timers';
 import { markModuleAndImpureDependenciesAsExecuted } from './utils/traverseStaticDependencies';
@@ -45,7 +46,6 @@ function normalizeEntryModules(
 export default class Graph {
 	acornParser: typeof acorn.Parser;
 	cachedModules: Map<string, ModuleJSON>;
-	contextParse: (code: string, acornOptions?: acorn.Options) => acorn.Node;
 	deoptimizationTracker: PathTracker;
 	entryModules: Module[] = [];
 	moduleLoader: ModuleLoader;
@@ -77,11 +77,6 @@ export default class Graph {
 				for (const key of Object.keys(cache)) cache[key][0]++;
 			}
 		}
-		this.contextParse = (code: string, options: Partial<acorn.Options> = {}) =>
-			this.acornParser.parse(code, {
-				...(this.options.acorn as acorn.Options),
-				...options
-			});
 
 		if (watcher) {
 			this.watchMode = true;
@@ -117,6 +112,35 @@ export default class Graph {
 		this.phase = BuildPhase.GENERATE;
 	}
 
+	contextParse(code: string, options: Partial<acorn.Options> = {}) {
+		const onCommentOrig = options.onComment;
+		const comments: acorn.Comment[] = [];
+
+		if (onCommentOrig && typeof onCommentOrig == 'function') {
+			options.onComment = (block, text, start, end, ...args) => {
+				comments.push({type: block ? "Block" : "Line", value: text, start, end});
+				return onCommentOrig.call(options, block, text, start, end, ...args);
+			}
+		} else {
+			options.onComment = comments;
+		}
+
+		const ast = this.acornParser.parse(code, {
+			...(this.options.acorn as acorn.Options),
+			...options
+		});
+
+		if (typeof onCommentOrig == 'object') {
+			onCommentOrig.push(...comments);
+		}
+
+		options.onComment = onCommentOrig;
+
+		markPureCallExpressions(comments, ast);
+
+		return ast;
+	}
+
 	getCache(): RollupCache {
 		// handle plugin cache eviction
 		for (const name in this.pluginCache) {

src/Module.ts

@@ -58,21 +58,13 @@ import { getOrCreate } from './utils/getOrCreate';
 import { getOriginalLocation } from './utils/getOriginalLocation';
 import { makeLegal } from './utils/identifierHelpers';
 import { basename, extname } from './utils/path';
-import { markPureCallExpressions } from './utils/pureComments';
 import relativeId from './utils/relativeId';
 import { RenderOptions } from './utils/renderHelpers';
 import { SOURCEMAPPING_URL_RE } from './utils/sourceMappingURL';
 import { timeEnd, timeStart } from './utils/timers';
 import { markModuleAndImpureDependenciesAsExecuted } from './utils/traverseStaticDependencies';
 import { MISSING_EXPORT_SHIM_VARIABLE } from './utils/variableNames';
 
-export interface CommentDescription {
-	block: boolean;
-	end: number;
-	start: number;
-	text: string;
-}
-
 interface ImportDescription {
 	module: Module | ExternalModule;
 	name: string;
@@ -122,35 +114,6 @@ export interface AstContext {
 	warn: (warning: RollupWarning, pos: number) => void;
 }
 
-function tryParse(
-	module: Module,
-	Parser: typeof acorn.Parser,
-	acornOptions: acorn.Options
-): acorn.Node {
-	try {
-		return Parser.parse(module.info.code!, {
-			...acornOptions,
-			onComment: (block: boolean, text: string, start: number, end: number) =>
-				module.comments.push({ block, text, start, end })
-		});
-	} catch (err) {
-		let message = err.message.replace(/ \(\d+:\d+\)$/, '');
-		if (module.id.endsWith('.json')) {
-			message += ' (Note that you need @rollup/plugin-json to import JSON files)';
-		} else if (!module.id.endsWith('.js')) {
-			message += ' (Note that you need plugins to import files that are not JavaScript)';
-		}
-		return module.error(
-			{
-				code: 'PARSE_ERROR',
-				message,
-				parserError: err
-			},
-			err.pos
-		);
-	}
-}
-
 const MISSING_EXPORT_SHIM_DESCRIPTION: ExportDescription = {
 	identifier: null,
 	localName: MISSING_EXPORT_SHIM_VARIABLE
@@ -218,7 +181,7 @@ export default class Module {
 	ast: Program | null = null;
 	chunkFileNames = new Set<string>();
 	chunkName: string | null = null;
-	comments: CommentDescription[] = [];
+	comments: acorn.Comment[] = [];
 	cycles = new Set<Symbol>();
 	dependencies = new Set<Module | ExternalModule>();
 	dynamicDependencies = new Set<Module | ExternalModule>();
@@ -719,13 +682,12 @@ export default class Module {
 
 		this.alwaysRemovedCode = alwaysRemovedCode || [];
 		if (!ast) {
-			ast = tryParse(this, this.graph.acornParser, this.options.acorn as acorn.Options);
+			ast = this.tryParse();
 			for (const comment of this.comments) {
-				if (!comment.block && SOURCEMAPPING_URL_RE.test(comment.text)) {
+				if (comment.type != "Block" && SOURCEMAPPING_URL_RE.test(comment.value)) {
 					this.alwaysRemovedCode.push([comment.start, comment.end]);
 				}
 			}
-			markPureCallExpressions(this.comments, ast);
 		}
 
 		timeEnd('generate ast', 3);
@@ -834,6 +796,32 @@ export default class Module {
 		return null;
 	}
 
+	tryParse(): acorn.Node {
+		try {
+			return this.graph.contextParse(this.info.code!, {
+				onComment: (block: boolean, text: string, start: number, end: number) =>
+					this.comments.push({ type: block ? "Block" : "Line", value: text, start, end })
+			});
+		} catch (err) {
+			let message = err.message.replace(/ \(\d+:\d+\)$/, '');
+			if (this.id.endsWith('.json')) {
+				message += ' (Note that you need @rollup/plugin-json to import JSON files)';
+			} else if (!this.id.endsWith('.js')) {
+				message += ' (Note that you need plugins to import files that are not JavaScript)';
+			}
+			console.log(err);
+			return this.error(
+				{
+					code: 'PARSE_ERROR',
+					message,
+					parserError: err
+				},
+				err.pos
+			);
+		}
+	}
+
+
 	updateOptions({
 		meta,
 		moduleSideEffects,

src/ast/keys.ts

@@ -9,7 +9,7 @@ export const keys: {
 
 export function getAndCreateKeys(esTreeNode: GenericEsTreeNode) {
 	keys[esTreeNode.type] = Object.keys(esTreeNode).filter(
-		key => typeof esTreeNode[key] === 'object'
+		key => key !== '_rollupAnnotations' && typeof esTreeNode[key] === 'object'
 	);
 	return keys[esTreeNode.type];
 }

src/ast/nodes/CallExpression.ts

@@ -21,12 +21,11 @@ import Identifier from './Identifier';
 import MemberExpression from './MemberExpression';
 import * as NodeType from './NodeType';
 import { ExpressionEntity } from './shared/Expression';
-import { ExpressionNode, IncludeChildren, INCLUDE_PARAMETERS, NodeBase } from './shared/Node';
+import { Annotation, ExpressionNode, IncludeChildren, INCLUDE_PARAMETERS, NodeBase } from './shared/Node';
 import SpreadElement from './SpreadElement';
 import Super from './Super';
 
 export default class CallExpression extends NodeBase implements DeoptimizableEntity {
-	annotatedPure?: boolean;
 	arguments!: (ExpressionNode | SpreadElement)[];
 	callee!: ExpressionNode | Super;
 	optional!: boolean;
@@ -157,7 +156,7 @@ export default class CallExpression extends NodeBase implements DeoptimizableEnt
 		}
 		if (
 			(this.context.options.treeshake as NormalizedTreeshakingOptions).annotations &&
-			this.annotatedPure
+			this.annotations?.some((a: Annotation) => a.pure)
 		)
 			return false;
 		return (

src/ast/nodes/NewExpression.ts

@@ -3,10 +3,9 @@ import { CallOptions } from '../CallOptions';
 import { HasEffectsContext } from '../ExecutionContext';
 import { EMPTY_PATH, ObjectPath, UNKNOWN_PATH } from '../utils/PathTracker';
 import * as NodeType from './NodeType';
-import { ExpressionNode, NodeBase } from './shared/Node';
+import { Annotation, ExpressionNode, NodeBase } from './shared/Node';
 
 export default class NewExpression extends NodeBase {
-	annotatedPure?: boolean;
 	arguments!: ExpressionNode[];
 	callee!: ExpressionNode;
 	type!: NodeType.tNewExpression;
@@ -27,7 +26,7 @@ export default class NewExpression extends NodeBase {
 		}
 		if (
 			(this.context.options.treeshake as NormalizedTreeshakingOptions).annotations &&
-			this.annotatedPure
+			this.annotations?.some((a: Annotation) => a.pure)
 		)
 			return false;
 		return (

src/ast/nodes/shared/Node.ts

@@ -1,7 +1,7 @@
 import * as acorn from 'acorn';
 import { locate } from 'locate-character';
 import MagicString from 'magic-string';
-import { AstContext, CommentDescription } from '../../../Module';
+import { AstContext } from '../../../Module';
 import { NodeRenderOptions, RenderOptions } from '../../../utils/renderHelpers';
 import { CallOptions } from '../../CallOptions';
 import { DeoptimizableEntity } from '../../DeoptimizableEntity';
@@ -26,9 +26,10 @@ export interface GenericEsTreeNode extends acorn.Node {
 
 export const INCLUDE_PARAMETERS: 'variables' = 'variables';
 export type IncludeChildren = boolean | typeof INCLUDE_PARAMETERS;
+export interface Annotation {comment?: acorn.Comment, pure?: boolean}
 
 export interface Node extends Entity {
-	annotations?: CommentDescription[];
+	annotations?: Annotation[];
 	context: AstContext;
 	end: number;
 	esTreeNode: GenericEsTreeNode;
@@ -88,6 +89,7 @@ export interface StatementNode extends Node {}
 export interface ExpressionNode extends ExpressionEntity, Node {}
 
 export class NodeBase implements ExpressionNode {
+	annotations?: Annotation[];
 	context: AstContext;
 	end!: number;
 	esTreeNode: acorn.Node;
@@ -126,7 +128,7 @@ export class NodeBase implements ExpressionNode {
 	bind() {
 		for (const key of this.keys) {
 			const value = (this as GenericEsTreeNode)[key];
-			if (value === null || key === 'annotations') continue;
+			if (value === null) continue;
 			if (Array.isArray(value)) {
 				for (const child of value) {
 					if (child !== null) child.bind();
@@ -165,7 +167,7 @@ export class NodeBase implements ExpressionNode {
 	hasEffects(context: HasEffectsContext): boolean {
 		for (const key of this.keys) {
 			const value = (this as GenericEsTreeNode)[key];
-			if (value === null || key === 'annotations') continue;
+			if (value === null) continue;
 			if (Array.isArray(value)) {
 				for (const child of value) {
 					if (child !== null && child.hasEffects(context)) return true;
@@ -195,7 +197,7 @@ export class NodeBase implements ExpressionNode {
 		this.included = true;
 		for (const key of this.keys) {
 			const value = (this as GenericEsTreeNode)[key];
-			if (value === null || key === 'annotations') continue;
+			if (value === null) continue;
 			if (Array.isArray(value)) {
 				for (const child of value) {
 					if (child !== null) child.include(context, includeChildrenRecursively);
@@ -232,7 +234,9 @@ export class NodeBase implements ExpressionNode {
 			// That way, we can override this function to add custom initialisation and then call super.parseNode
 			if (this.hasOwnProperty(key)) continue;
 			const value = esTreeNode[key];
-			if (typeof value !== 'object' || value === null || key === 'annotations') {
+			if (key === '_rollupAnnotations') {
+				this.annotations = value;
+			} else if (typeof value !== 'object' || value === null) {
 				(this as GenericEsTreeNode)[key] = value;
 			} else if (Array.isArray(value)) {
 				(this as GenericEsTreeNode)[key] = [];
@@ -254,7 +258,7 @@ export class NodeBase implements ExpressionNode {
 	render(code: MagicString, options: RenderOptions) {
 		for (const key of this.keys) {
 			const value = (this as GenericEsTreeNode)[key];
-			if (value === null || key === 'annotations') continue;
+			if (value === null) continue;
 			if (Array.isArray(value)) {
 				for (const child of value) {
 					if (child !== null) child.render(code, options);

src/utils/PluginContext.ts

@@ -156,7 +156,7 @@ export function getPluginContexts(
 				const moduleIds = graph.modulesById.keys();
 				return wrappedModuleIds();
 			},
-			parse: graph.contextParse,
+			parse: graph.contextParse.bind(graph),
 			resolve(source, importer, { custom, skipSelf } = BLANK) {
 				return graph.moduleLoader.resolveId(source, importer, custom, skipSelf ? pidx : null);
 			},

src/utils/pureComments.ts

@@ -1,7 +1,7 @@
 import * as acorn from 'acorn';
 import { base as basicWalker, BaseWalker } from 'acorn-walk';
 import { CallExpression, ExpressionStatement, NewExpression } from '../ast/nodes/NodeType';
-import { CommentDescription } from '../Module';
+import { Annotation } from '../ast/nodes/shared/Node';
 
 // patch up acorn-walk until class-fields are officially supported
 basicWalker.PropertyDefinition = function (node: any, st: any, c: any) {
@@ -15,7 +15,7 @@ basicWalker.PropertyDefinition = function (node: any, st: any, c: any) {
 
 interface CommentState {
 	commentIndex: number;
-	commentNodes: CommentDescription[];
+	commentNodes: acorn.Comment[];
 }
 
 function handlePureAnnotationsOfNode(
@@ -34,26 +34,30 @@ function handlePureAnnotationsOfNode(
 }
 
 function markPureNode(
-	node: acorn.Node & { annotations?: CommentDescription[] },
-	comment: CommentDescription
+	node: acorn.Node & { _rollupAnnotations?: Annotation[] },
+	comment: acorn.Comment
 ) {
-	if (node.annotations) {
-		node.annotations.push(comment);
+	if (node._rollupAnnotations) {
+		node._rollupAnnotations.push({comment});
 	} else {
-		node.annotations = [comment];
+		node._rollupAnnotations = [{comment}];
 	}
 	if (node.type === ExpressionStatement) {
 		node = (node as any).expression;
 	}
 	if (node.type === CallExpression || node.type === NewExpression) {
-		(node as any).annotatedPure = true;
+		if (node._rollupAnnotations) {
+			node._rollupAnnotations.push({pure: true});
+		} else {
+			node._rollupAnnotations = [{pure: true}];
+		}
 	}
 }
 
 const pureCommentRegex = /[@#]__PURE__/;
-const isPureComment = (comment: CommentDescription) => pureCommentRegex.test(comment.text);
+const isPureComment = (comment: acorn.Comment) => pureCommentRegex.test(comment.value);
 
-export function markPureCallExpressions(comments: CommentDescription[], esTreeAst: acorn.Node) {
+export function markPureCallExpressions(comments: acorn.Comment[], esTreeAst: acorn.Node) {
 	handlePureAnnotationsOfNode(esTreeAst, {
 		commentIndex: 0,
 		commentNodes: comments.filter(isPureComment)

src/utils/treeshakeNode.ts

@@ -6,8 +6,11 @@ export function treeshakeNode(node: Node, code: MagicString, start: number, end:
 	code.remove(start, end);
 	if (node.annotations) {
 		for (const annotation of node.annotations) {
-			if (annotation.start < start) {
-				code.remove(annotation.start, annotation.end);
+			if (!annotation.comment) {
+				continue;
+			}
+			if (annotation.comment.start < start) {
+				code.remove(annotation.comment.start, annotation.comment.end);
 			} else {
 				return;
 			}
@@ -20,8 +23,8 @@ export function removeAnnotations(node: Node, code: MagicString) {
 		node = node.parent as Node;
 	}
 	if (node.annotations) {
-		for (const annotation of node.annotations) {
-			code.remove(annotation.start, annotation.end);
+		for (const annotation of node.annotations.filter((a) => a.comment)) {
+			code.remove(annotation.comment!.start, annotation.comment!.end);
 		}
 	}
 }