Review comments: Extract and stream-line code, add edge-case test

src/Module.ts

@@ -1,5 +1,5 @@
 import { IParse, Options as AcornOptions } from 'acorn';
-import walk from 'acorn/dist/walk';
+// @ts-ignore
 import * as ESTree from 'estree';
 import { locate } from 'locate-character';
 import MagicString from 'magic-string';
@@ -36,6 +36,7 @@ import getCodeFrame from './utils/getCodeFrame';
 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';
@@ -160,38 +161,6 @@ function handleMissingExport(
 	);
 }
 
-// Find syntax nodes following comments
-function findNodesAfterComments(ast, commentNodes) {
-	const state = {
-		commentIdx: 0,
-		nodes: []
-	};
-
-	(function c(node, state, override) {
-		if (state.commentIdx === commentNodes.length) return;
-		const pos = commentNodes[state.commentIdx].end;
-		if (node.end < pos) return;
-		const type = override || node.type;
-		if (node.start >= pos) {
-			state.commentIdx++;
-			state.nodes.push(node);
-		}
-		walk.base[type](node, state, c);
-	})(ast, state);
-
-	return state.nodes;
-}
-
-function markNodePure(node) {
-	if (node.type === 'ExpressionStatement') {
-		markNodePure(node.expression);
-	} else if (node.type === 'CallExpression') {
-		node.markedPure = true;
-	}
-}
-
-const pureCommentRegex = /^ ?#__PURE__\s*$/;
-
 export default class Module {
 	type: 'Module';
 	private graph: Graph;
@@ -288,14 +257,10 @@ export default class Module {
 		timeStart('generate ast', 3);
 
 		this.esTreeAst = ast || tryParse(this, this.graph.acornParse, this.graph.acornOptions);
+		markPureCallExpressions(this.comments, this.esTreeAst);
 
 		timeEnd('generate ast', 3);
 
-		const pureMarkerComments = this.comments.filter(comment => pureCommentRegex.test(comment.text));
-		const nodesAfterPureComments = findNodesAfterComments(this.esTreeAst, pureMarkerComments);
-
-		nodesAfterPureComments.forEach(node => markNodePure(node));
-
 		this.resolvedIds = resolvedIds || Object.create(null);
 
 		// By default, `id` is the filename. Custom resolvers and loaders

src/ast/nodes/CallExpression.ts

@@ -23,6 +23,7 @@ export default class CallExpression extends NodeBase implements DeoptimizableEnt
 	type: NodeType.tCallExpression;
 	callee: ExpressionNode;
 	arguments: (ExpressionNode | SpreadElement)[];
+	pure?: boolean;
 
 	private callOptions: CallOptions;
 
@@ -135,7 +136,7 @@ export default class CallExpression extends NodeBase implements DeoptimizableEnt
 		for (const argument of this.arguments) {
 			if (argument.hasEffects(options)) return true;
 		}
-		if (this.markedPure) return false;
+		if (this.pure) return false;
 		return (
 			this.callee.hasEffects(options) ||
 			this.callee.hasEffectsWhenCalledAtPath(

src/ast/nodes/ExpressionStatement.ts

@@ -1,10 +1,11 @@
 import MagicString from 'magic-string';
 import { RenderOptions } from '../../utils/renderHelpers';
 import * as NodeType from './NodeType';
-import { StatementBase } from './shared/Node';
+import { ExpressionNode, StatementBase } from './shared/Node';
 
 export default class ExpressionStatement extends StatementBase {
 	directive?: string;
+	expression: ExpressionNode;
 
 	initialise() {
 		this.included = false;

src/ast/nodes/Identifier.ts

@@ -10,7 +10,6 @@ import { ImmutableEntityPathTracker } from '../utils/ImmutableEntityPathTracker'
 import { LiteralValueOrUnknown, ObjectPath, UNKNOWN_EXPRESSION, UNKNOWN_VALUE } from '../values';
 import LocalVariable from '../variables/LocalVariable';
 import Variable from '../variables/Variable';
-import AssignmentExpression from './AssignmentExpression';
 import AssignmentPattern from './AssignmentPattern';
 import * as NodeType from './NodeType';
 import Property from './Property';

src/ast/nodes/shared/Node.ts

@@ -91,7 +91,6 @@ export class NodeBase implements ExpressionNode {
 
 	constructor(
 		esTreeNode: GenericEsTreeNode,
-		// we need to pass down the node constructors to avoid a circular dependency
 		parent: Node | { type: string; context: AstContext },
 		parentScope: Scope
 	) {

src/utils/pureComments.ts

@@ -0,0 +1,46 @@
+// @ts-ignore
+import walk from 'acorn/dist/walk';
+import * as ESTree from 'estree';
+import { CommentDescription } from '../Module';
+
+function forEachNodeAfterComment(
+	ast: ESTree.Program,
+	commentNodes: CommentDescription[],
+	handleNode: (node: ESTree.Node) => void
+): void {
+	let commentIdx = 0;
+
+	checkCommentsBeforeNode(<any>ast);
+
+	function checkCommentsBeforeNode(node: ESTree.Node & { start: number; end: number }) {
+		if (commentIdx === commentNodes.length || commentNodes[commentIdx].end > node.end) return;
+		let isFound = false;
+		while (commentIdx < commentNodes.length && node.start >= commentNodes[commentIdx].end) {
+			if (!isFound) {
+				handleNode(node);
+				isFound = true;
+			}
+			commentIdx++;
+		}
+		walk.base[node.type](node, null, checkCommentsBeforeNode);
+	}
+}
+
+function markPureNode(node: ESTree.Node) {
+	while (node.type === 'ExpressionStatement') {
+		node = node.expression;
+	}
+	if (node.type === 'CallExpression') {
+		(<any>node).pure = true;
+	}
+}
+
+const pureCommentRegex = /^ ?#__PURE__\s*$/;
+
+export function markPureCallExpressions(comments: CommentDescription[], esTreeAst: ESTree.Program) {
+	forEachNodeAfterComment(
+		esTreeAst,
+		comments.filter(comment => pureCommentRegex.test(comment.text)),
+		markPureNode
+	);
+}

test/form/samples/pure-comments-multiple/_config.js

@@ -0,0 +1,3 @@
+module.exports = {
+	description: 'does not associate multiple "pure" comments before a token with subsequent tokens'
+};

test/form/samples/pure-comments-multiple/_expected.js

@@ -0,0 +1 @@
+retained();

test/form/samples/pure-comments-multiple/main.js

@@ -0,0 +1,2 @@
+/*#__PURE__*//*#__PURE__*//*#__PURE__*//*#__PURE__*//*#__PURE__*/ removed();
+retained();

test/function/samples/call-marked-pure/_config.js

@@ -1,26 +1,20 @@
 const assert = require('assert');
 
 module.exports = {
-	description: 'functions marked with pure comment do not have effects',
+	description: 'external function calls marked with pure comment do not have effects',
+	options: {
+		external: ['socks']
+	},
 	context: {
 		require(id) {
 			if (id === 'socks') {
-				return () => '🧦';
+				return () => {
+					throw new Error('Not all socks were removed.');
+				};
 			}
 		}
 	},
 	code(code) {
 		assert.ok(code.search(/socks\(\)/) === -1);
-	},
-	warnings: [
-		{
-			code: 'UNRESOLVED_IMPORT',
-			importer: 'main.js',
-			message:
-				"'socks' is imported by main.js, but could not be resolved – treating it as an external dependency",
-			source: 'socks',
-			url:
-				'https://github.com/rollup/rollup/wiki/Troubleshooting#treating-module-as-external-dependency'
-		}
-	]
+	}
 };

test/function/samples/call-marked-pure/main.js

@@ -4,5 +4,5 @@ import socks from 'socks';
 /* #__PURE__*/ socks();
 /*#__PURE__ */ socks();
 /* #__PURE__ */ socks();
-// #__PURE__        
-socks();
+// #__PURE__
+socks();