Correctly associate pure annotations and remove invalid ones (#4095)

* Simplify annotation handling

* Remove invalid annotations and improve annotation association

* Improve coverage and remove trailing annotations

src/Graph.ts

@@ -17,7 +17,7 @@ import { PluginDriver } from './utils/PluginDriver';
 import { BuildPhase } from './utils/buildPhase';
 import { errImplicitDependantIsNotIncluded, error } from './utils/error';
 import { analyseModuleExecution } from './utils/executionOrder';
-import { markPureCallExpressions } from './utils/pureComments';
+import { addAnnotations } from './utils/pureComments';
 import relativeId from './utils/relativeId';
 import { timeEnd, timeStart } from './utils/timers';
 import { markModuleAndImpureDependenciesAsExecuted } from './utils/traverseStaticDependencies';
@@ -137,7 +137,7 @@ export default class Graph {
 
 		options.onComment = onCommentOrig;
 
-		markPureCallExpressions(comments, ast, code);
+		addAnnotations(comments, ast, code);
 
 		return ast;
 	}

src/Module.ts

@@ -19,7 +19,7 @@ import Program from './ast/nodes/Program';
 import TemplateLiteral from './ast/nodes/TemplateLiteral';
 import VariableDeclaration from './ast/nodes/VariableDeclaration';
 import { nodeConstructors } from './ast/nodes/index';
-import { ExpressionNode, GenericEsTreeNode, NodeBase } from './ast/nodes/shared/Node';
+import { ExpressionNode, NodeBase } from './ast/nodes/shared/Node';
 import ModuleScope from './ast/scopes/ModuleScope';
 import { PathTracker, UNKNOWN_PATH } from './ast/utils/PathTracker';
 import ExportDefaultVariable from './ast/variables/ExportDefaultVariable';
@@ -62,7 +62,6 @@ import { makeLegal } from './utils/identifierHelpers';
 import { basename, extname } from './utils/path';
 import relativeId from './utils/relativeId';
 import { RenderOptions } from './utils/renderHelpers';
-import { SOURCEMAPPING_URL_COMMENT_RE } from './utils/sourceMappingURL';
 import { timeEnd, timeStart } from './utils/timers';
 import { markModuleAndImpureDependenciesAsExecuted } from './utils/traverseStaticDependencies';
 import { MISSING_EXPORT_SHIM_VARIABLE } from './utils/variableNames';
@@ -121,34 +120,6 @@ const MISSING_EXPORT_SHIM_DESCRIPTION: ExportDescription = {
 	localName: MISSING_EXPORT_SHIM_VARIABLE
 };
 
-function findSourceMappingURLComments(ast: acorn.Node, code: string): [number, number][] {
-	const ret: [number, number][] = [];
-
-	const addCommentsPos = (start: number, end: number): void => {
-		if (start == end) {
-			return;
-		}
-
-		let sourcemappingUrlMatch;
-		const interStatmentCode = code.slice(start, end);
-		while ((sourcemappingUrlMatch = SOURCEMAPPING_URL_COMMENT_RE.exec(interStatmentCode))) {
-			ret.push([
-				start + sourcemappingUrlMatch.index,
-				start + SOURCEMAPPING_URL_COMMENT_RE.lastIndex
-			]);
-		}
-	};
-
-	let prevStmtEnd = 0;
-	for (const stmt of (ast as GenericEsTreeNode).body) {
-		addCommentsPos(prevStmtEnd, stmt.start);
-		prevStmtEnd = stmt.end;
-	}
-	addCommentsPos(prevStmtEnd, code.length);
-
-	return ret;
-}
-
 function getVariableForExportNameRecursive(
 	target: Module | ExternalModule,
 	name: string,
@@ -251,7 +222,6 @@ export default class Module {
 	usesTopLevelAwait = false;
 
 	private allExportNames: Set<string> | null = null;
-	private alwaysRemovedCode!: [number, number][];
 	private astContext!: AstContext;
 	private readonly context: string;
 	private customTransformCache!: boolean;
@@ -688,7 +658,6 @@ export default class Module {
 	}
 
 	setSource({
-		alwaysRemovedCode,
 		ast,
 		code,
 		customTransformCache,
@@ -700,7 +669,6 @@ export default class Module {
 		transformFiles,
 		...moduleOptions
 	}: TransformModuleJSON & {
-		alwaysRemovedCode?: [number, number][];
 		transformFiles?: EmittedFile[] | undefined;
 	}): void {
 		this.info.code = code;
@@ -716,11 +684,9 @@ export default class Module {
 
 		timeStart('generate ast', 3);
 
-		this.alwaysRemovedCode = alwaysRemovedCode || [];
 		if (!ast) {
 			ast = this.tryParse();
 		}
-		this.alwaysRemovedCode.push(...findSourceMappingURLComments(ast, this.info.code));
 
 		timeEnd('generate ast', 3);
 
@@ -734,9 +700,6 @@ export default class Module {
 			filename: (this.excludeFromSourcemap ? null : fileName)!, // don't include plugin helpers in sourcemap
 			indentExclusionRanges: []
 		});
-		for (const [start, end] of this.alwaysRemovedCode) {
-			this.magicString.remove(start, end);
-		}
 
 		timeStart('analyse ast', 3);
 
@@ -778,7 +741,6 @@ export default class Module {
 
 	toJSON(): ModuleJSON {
 		return {
-			alwaysRemovedCode: this.alwaysRemovedCode,
 			ast: this.ast!.esTreeNode,
 			code: this.info.code!,
 			customTransformCache: this.customTransformCache,

src/ast/keys.ts

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

src/ast/nodes/CallExpression.ts

@@ -28,13 +28,7 @@ import {
 	UNKNOWN_EXPRESSION,
 	UnknownValue
 } from './shared/Expression';
-import {
-	Annotation,
-	ExpressionNode,
-	INCLUDE_PARAMETERS,
-	IncludeChildren,
-	NodeBase
-} from './shared/Node';
+import { ExpressionNode, INCLUDE_PARAMETERS, IncludeChildren, NodeBase } from './shared/Node';
 
 export default class CallExpression extends NodeBase implements DeoptimizableEntity {
 	arguments!: (ExpressionNode | SpreadElement)[];
@@ -188,7 +182,7 @@ export default class CallExpression extends NodeBase implements DeoptimizableEnt
 		}
 		if (
 			(this.context.options.treeshake as NormalizedTreeshakingOptions).annotations &&
-			this.annotations?.some((a: Annotation) => a.pure)
+			this.annotations
 		)
 			return false;
 		return (

src/ast/nodes/IfStatement.ts

@@ -1,6 +1,5 @@
 import MagicString from 'magic-string';
 import { RenderOptions } from '../../utils/renderHelpers';
-import { removeAnnotations } from '../../utils/treeshakeNode';
 import { DeoptimizableEntity } from '../DeoptimizableEntity';
 import { BROKEN_FLOW_NONE, HasEffectsContext, InclusionContext } from '../ExecutionContext';
 import TrackingScope from '../scopes/TrackingScope';
@@ -89,7 +88,6 @@ export default class IfStatement extends StatementBase implements DeoptimizableE
 		if (includesIfElse) {
 			this.test.render(code, options);
 		} else {
-			removeAnnotations(this, code);
 			code.remove(this.start, this.consequent.start);
 		}
 		if (this.consequent.included && (noTreeshake || testValue === UnknownValue || testValue)) {

src/ast/nodes/NewExpression.ts

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

src/ast/nodes/shared/Node.ts

@@ -2,6 +2,7 @@ import * as acorn from 'acorn';
 import { locate, Location } from 'locate-character';
 import MagicString from 'magic-string';
 import { AstContext } from '../../../Module';
+import { ANNOTATION_KEY, INVALID_COMMENT_KEY } from '../../../utils/pureComments';
 import { NodeRenderOptions, RenderOptions } from '../../../utils/renderHelpers';
 import { Entity } from '../../Entity';
 import {
@@ -21,13 +22,9 @@ export interface GenericEsTreeNode extends acorn.Node {
 
 export const INCLUDE_PARAMETERS = 'variables' as const;
 export type IncludeChildren = boolean | typeof INCLUDE_PARAMETERS;
-export interface Annotation {
-	comment?: acorn.Comment;
-	pure?: boolean;
-}
 
 export interface Node extends Entity {
-	annotations?: Annotation[];
+	annotations?: acorn.Comment[];
 	context: AstContext;
 	end: number;
 	esTreeNode: GenericEsTreeNode;
@@ -87,7 +84,7 @@ export type StatementNode = Node;
 export interface ExpressionNode extends ExpressionEntity, Node {}
 
 export class NodeBase extends ExpressionEntity implements ExpressionNode {
-	annotations?: Annotation[];
+	annotations?: acorn.Comment[];
 	context: AstContext;
 	end!: number;
 	esTreeNode: acorn.Node;
@@ -201,8 +198,13 @@ export class NodeBase extends ExpressionEntity implements ExpressionNode {
 		for (const [key, value] of Object.entries(esTreeNode)) {
 			// That way, we can override this function to add custom initialisation and then call super.parseNode
 			if (this.hasOwnProperty(key)) continue;
-			if (key === '_rollupAnnotations') {
-				this.annotations = value;
+			if (key.charCodeAt(0) === 95 /* _ */) {
+				if (key === ANNOTATION_KEY) {
+					this.annotations = value;
+				} else if (key === INVALID_COMMENT_KEY) {
+					for (const { start, end } of value as acorn.Comment[])
+						this.context.magicString.remove(start, end);
+				}
 			} else if (typeof value !== 'object' || value === null) {
 				(this as GenericEsTreeNode)[key] = value;
 			} else if (Array.isArray(value)) {

src/rollup/types.d.ts

@@ -114,7 +114,6 @@ export interface TransformModuleJSON extends Partial<PartialNull<ModuleOptions>>
 }
 
 export interface ModuleJSON extends TransformModuleJSON {
-	alwaysRemovedCode: [number, number][];
 	ast: AcornNode;
 	dependencies: string[];
 	id: string;