Improve AST checking and types (#189)

Author: dario-piotrowicz

.changeset/short-goats-grab.md

@@ -0,0 +1,8 @@
+---
+'@cloudflare/next-on-pages': patch
+---
+
+improve AST checking
+
+improve the way we check for webpack chunks (for the experimental minification) by
+improving the AST types used and also make the AST checking more robust

package-lock.json

@@ -22,6 +22,7 @@
 	],
 	"dependencies": {
 		"acorn": "^8.8.0",
+		"ast-types": "^0.14.2",
 		"astring": "^1.8.4",
 		"chalk": "^5.2.0",
 		"chokidar": "^3.5.3",

package.json

@@ -1,7 +1,6 @@
 import { readFile, writeFile, mkdir, rm, readdir, copyFile } from 'fs/promises';
 import { exit } from 'process';
 import { dirname, join, relative, resolve } from 'path';
-import type { Node } from 'acorn';
 import { parse } from 'acorn';
 import { generate } from 'astring';
 import {
@@ -15,6 +14,8 @@ import {
 import type { CliOptions } from '../cli';
 import { cliError, cliWarn } from '../cli';
 import { tmpdir } from 'os';
+import type * as AST from 'ast-types/gen/kinds';
+import assert from 'node:assert';
 
 /**
  * Creates new files containing the Vercel built functions but adjusted so that they can be later
@@ -283,69 +284,53 @@ function extractWebpackChunks(
 	const parsedContents = parse(functionContents, {
 		ecmaVersion: 'latest',
 		sourceType: 'module',
-	}) as Node & { body: LooseNode[] };
-
-	const expressions = parsedContents.body
-		.filter(
-			({ type, expression }) =>
-				type === 'ExpressionStatement' &&
-				expression?.type === 'CallExpression' &&
-				expression.callee?.type === 'MemberExpression' &&
-				expression.callee.object?.type === 'AssignmentExpression' &&
-				expression.callee.object.left?.object?.name === 'self' &&
-				expression.callee.object.left.property?.name === 'webpackChunk_N_E' &&
-				expression.arguments?.[0]?.elements?.[1]?.type === 'ObjectExpression'
-		)
-		.map(
-			node => node?.expression?.arguments?.[0]?.elements?.[1]?.properties
-		) as LooseNode[][];
-
-	for (const objectOfChunks of expressions) {
-		for (const chunkExpression of objectOfChunks) {
-			const key = chunkExpression?.key?.value;
-			if (key in existingWebpackChunks) {
-				if (
-					existingWebpackChunks.get(key) !== generate(chunkExpression.value)
-				) {
-					cliError(
-						`
+	}) as unknown as AST.ProgramKind;
+
+	const chunks = parsedContents.body.flatMap(getWebpackChunksFromStatement);
+
+	for (const chunk of chunks) {
+		const key = (chunk.key as AST.NumericLiteralKind).value;
+
+		if (key in existingWebpackChunks) {
+			if (existingWebpackChunks.get(key) !== generate(chunk.value)) {
+				cliError(
+					`
 							ERROR: Detected a collision with '--experimental-minify'.
 							       Try removing the '--experimental-minify' argument.
 						`,
-						{ spaced: true }
-					);
-					exit(1);
-				}
+					{ spaced: true }
+				);
+				exit(1);
 			}
+		}
 
-			webpackChunks.set(key, generate(chunkExpression.value));
+		webpackChunks.set(key, generate(chunk.value));
 
-			const chunkFilePath = join(tmpWebpackDir, `${key}.js`);
+		const chunkFilePath = join(tmpWebpackDir, `${key}.js`);
 
-			const newValue = {
-				type: 'MemberExpression',
-				object: {
-					type: 'CallExpression',
-					callee: {
-						type: 'Identifier',
-						name: 'require',
-					},
-					arguments: [
-						{
-							type: 'Literal',
-							value: chunkFilePath,
-							raw: JSON.stringify(chunkFilePath),
-						},
-					],
-				},
-				property: {
+		const newValue = {
+			type: 'MemberExpression',
+			object: {
+				type: 'CallExpression',
+				callee: {
 					type: 'Identifier',
-					name: 'default',
+					name: 'require',
 				},
-			};
+				arguments: [
+					{
+						type: 'Literal',
+						value: chunkFilePath,
+						raw: JSON.stringify(chunkFilePath),
+					},
+				],
+			},
+			property: {
+				type: 'Identifier',
+				name: 'default',
+			},
+		};
 
-			chunkExpression.value = newValue;
-		}
+		(chunk as unknown as { value: unknown }).value = newValue;
 	}
 
 	return {
@@ -418,21 +403,71 @@ type DirectoryProcessingResults = {
 	webpackChunks: Map<number, string>;
 };
 
-type LooseNode = Node & {
-	expression?: LooseNode;
-	callee?: LooseNode;
-	object?: LooseNode;
-	left?: LooseNode;
-	right?: LooseNode;
-	property?: LooseNode;
-	arguments?: LooseNode[];
-	elements?: LooseNode[];
-	properties?: LooseNode[];
-	key?: LooseNode;
-	name?: string;
-	/*
-    eslint-disable-next-line @typescript-eslint/no-explicit-any 
-    -- TODO: improve the type of value
-  */
-	value: any;
-};
+/**
+ * Verifies wether the provided AST statement represents a javascript code
+ * of the following format:
+ * 	```
+ * 		(self.webpackChunk_N_E = self.webpackChunk_N_E || []).push(...,
+ * 			{
+ * 				[chunkNumberA]: e => { ... },
+ * 				[chunkNumberB]: e => { ... },
+ * 				[chunkNumberC]: e => { ... },
+ * 				...
+ * 			}
+ * 		]);
+ *  ```
+ * and in such case it extracts the various chunk properties.
+ *
+ * @param statement the AST statement to check
+ *
+ * @returns the chunks as an array of AST Properties if the statement represent the target javascript code, an empty array otherwise
+ */
+function getWebpackChunksFromStatement(
+	statement: AST.StatementKind
+): AST.PropertyKind[] {
+	try {
+		assert(statement.type === 'ExpressionStatement');
+		const expr = statement.expression;
+
+		assert(expr.type === 'CallExpression');
+		assert(expr.callee.type === 'MemberExpression');
+		assert(expr.callee.property.type === 'Identifier');
+		assert(expr.callee.property.name === 'push');
+		const calleeObj = expr.callee.object;
+
+		assert(calleeObj.type === 'AssignmentExpression');
+
+		assertSelfWebpackChunk_N_E(calleeObj.left);
+
+		assert(calleeObj.right.type === 'LogicalExpression');
+		assert(calleeObj.right.operator === '||');
+		assertSelfWebpackChunk_N_E(calleeObj.right.left);
+		assert(calleeObj.right.right.type === 'ArrayExpression');
+		assert(calleeObj.right.right.elements.length === 0);
+
+		assert(expr.arguments[0]?.type === 'ArrayExpression');
+		assert(expr.arguments[0].elements[1]?.type === 'ObjectExpression');
+
+		return expr.arguments[0].elements[1].properties.filter(
+			p =>
+				p.type === 'Property' &&
+				p.key.type === 'Literal' &&
+				typeof p.key.value === 'number' &&
+				p.value.type === 'ArrowFunctionExpression'
+		) as AST.PropertyKind[];
+	} catch {
+		return [];
+	}
+}
+
+/**
+ * Asserts whether the provided AST node represents `self.webpackChunk_N_E`
+ * (throws an AssertionError it doesn't)
+ */
+function assertSelfWebpackChunk_N_E(expression: AST.NodeKind): void {
+	assert(expression.type === 'MemberExpression');
+	assert(expression.object.type === 'Identifier');
+	assert(expression.object.name === 'self');
+	assert(expression.property.type === 'Identifier');
+	assert(expression.property.name === 'webpackChunk_N_E');
+}