1.5.0 Optimization and Async (#22)

* Refactor generation of flat AST and optimize scoping and attaching references.
Fix issue with scope blocks appearing as their own scope. Function now appear in their upper scope, rather than in themselves.
 Create generateFlastASTAsync. It's a bit slower than the non-async version, but at least it doesn't block.

* Add missing properties

* 1.5.0

package.json

@@ -1,6 +1,6 @@
 {
   "name": "flast",
-  "version": "1.4.0",
+  "version": "1.5.0",
   "description": "Flatten JS AST",
   "main": "src/index.js",
   "scripts": {

src/flast.js

@@ -72,9 +72,9 @@ function generateFlatAST(inputCode, opts = {}) {
 	opts = { ...generateFlatASTDefaultOptions, ...opts };
 	const rootNode = generateRootNode(inputCode, opts);
 	const tree = extractNodesFromRoot(rootNode, opts);
-	const sm = initScopeManager(rootNode);
 	if (opts.detailed) {
-		for (let i = 0; i < tree.length; i++) injectScopeToNode(tree[i], sm);
+		const scopes = getAllScopes(rootNode);
+		for (let i = 0; i < tree.length; i++) injectScopeToNode(tree[i], scopes);
 	}
 	return tree;
 }
@@ -119,6 +119,7 @@ function extractNodesFromRoot(rootNode, opts) {
 	const tree = [];
 	let nodeId = 0;
 
+	// noinspection JSUnusedGlobalSymbols
 	estraverse.traverse(rootNode, {
 		/**
 		 * @param {ASTNode} node
@@ -129,8 +130,12 @@ function extractNodesFromRoot(rootNode, opts) {
 			node.nodeId = nodeId++;
 			node.childNodes = [];
 			node.parentNode = parentNode;
-			// Keep track of the node's lineage
 			node.parentKey = parentNode ? getParentKey(node) : '';
+			node.lineage = [...parentNode?.lineage || []];
+			if (parentNode) {
+				node.lineage.push(parentNode.nodeId);
+				parentNode.childNodes.push(node);
+			}
 			if (opts.includeSrc) Object.defineProperty(node, 'src', {
 				get() { return rootNode.srcClosure(node.range[0], node.range[1]);},
 			});
@@ -139,71 +144,147 @@ function extractNodesFromRoot(rootNode, opts) {
 	return tree;
 }
 
-function initScopeManager(rootNode) {
-	// noinspection JSCheckFunctionSignatures
-	return analyze(rootNode, {
-		optimistic: true,
-		ecmaVersion,
-		sourceType});
-}
-
 /**
- *
  * @param {ASTNode} node
- * @param {ScopeManager} sm
+ * @param {ASTScope[]} scopes
  */
-function injectScopeToNode(node, sm) {
+function injectScopeToNode(node, scopes) {
 	let parentNode = node.parentNode;
 	// Acquire scope
-	node.scope = sm.acquire(node);
-	if (!node.scope) node.scope = node.parentNode.scope;
-	else if (node.scope.type.includes('-name') && node.scope?.childScopes?.length === 1) node.scope = node.scope.childScopes[0];
-	if (node.scope.scopeId === undefined) node.scope.scopeId = node.scope.block.nodeId;
-	if (parentNode) {
-		node.lineage = [...parentNode?.lineage || [], parentNode.nodeId];
-		parentNode.childNodes.push(node);
-	}
-	if (node.type === 'Identifier') {
+	node.scope = matchScopeToNode(node, scopes);
+	if (node.type === 'Identifier' && !(!parentNode.computed && ['property', 'key'].includes(node.parentKey))) {
 		// Track references and declarations
 		// Prevent assigning declNode to member expression properties or object keys
-		if (!(['property', 'key'].includes(node.parentKey) && !parentNode.computed)) {
-			const variables = node.scope.variables.filter(n => n.name === node.name);
-			const isDeclaration = variables?.length && variables[0].identifiers.filter(n => n.nodeId === node.nodeId).length;
-			if (isDeclaration) node.references = node.references || [];
-			else if (!(node.parentKey === 'id' && node.parentNode.type === 'FunctionDeclaration'))  {
-				// Find declaration by finding the closest declaration of the same name.
-				let decls = [];
-				if (variables?.length) decls = variables.filter(n => n.name === node.name)[0].identifiers;
-				else {
-					const scopeReferences = node.scope.references.filter(n => n.identifier.name === node.name);
-					if (scopeReferences.length) decls = scopeReferences[0].resolved?.identifiers || [];
-				}
-				let declNode = decls[0];
-				if (decls.length > 1) {   // TODO: Defer setting declaration and references
-					let commonAncestors = node.lineage.reduce((t, c) => declNode.lineage?.includes(c) ? ++t : t, 0);
-					decls.slice(1).forEach(n => {
-						const ca = node.lineage.reduce((t, c) => n.lineage?.includes(c) ? ++t : t, 0);
-						if (ca > commonAncestors) {
-							commonAncestors = ca;
-							declNode = n;
-						}
-					});
-				}
-				if (declNode) {
-					if (!declNode.references) declNode.references = [];
-					declNode.references.push(node);
-					node.declNode = declNode;
+		const variables = node.scope.variables.filter(n => n.name === node.name);
+		if (node.parentKey === 'id' || (variables?.length && variables[0].identifiers.some(n => n === node))) {
+			node.references = node.references || [];
+		} else {
+			// Find declaration by finding the closest declaration of the same name.
+			let decls = [];
+			if (variables?.length) {
+				decls = variables.find(n => n.name === node.name)?.identifiers;
+			}
+			else {
+				const scopeReference = node.scope.references.find(n => n.identifier.name === node.name);
+				if (scopeReference) decls = scopeReference.resolved?.identifiers || [];
+			}
+			let declNode = decls[0];
+			if (decls.length > 1) {
+				let commonAncestors = maxSharedLength(declNode.lineage, node.lineage);
+				for (let i = 1; i < decls.length; i++) {
+					const ca = maxSharedLength(decls[i].lineage, node.lineage);
+					if (ca > commonAncestors) {
+						commonAncestors = ca;
+						declNode = decls[i];
+					}
 				}
 			}
+			if (declNode) {
+				declNode.references = declNode.references || [];
+				declNode.references.push(node);
+				node.declNode = declNode;
+			}
+		}
+	}
+}
+
+/**
+ * @param {number[]} targetArr
+ * @param {number[]} containedArr
+ * @return {number} Return the maximum length of shared numbers
+ */
+function maxSharedLength(targetArr, containedArr) {
+	let count = 0;
+	for (let i = 0; i < containedArr.length; i++) {
+		if (targetArr[i] !== containedArr[i]) break;
+		++count;
+	}
+	return count;
+}
+
+/**
+ * @param {ASTNode} node
+ * @param {ASTScope[]} scopes
+ * @return {Promise}
+ */
+async function injectScopeToNodeAsync(node, scopes) {
+	return new Promise((resolve, reject) => {
+		try {
+			injectScopeToNode(node, scopes);
+			resolve();
+		} catch (e) {
+			reject(e);
 		}
+	});
+}
+
+function getAllScopes(rootNode) {
+	const globalScope = analyze(rootNode, {
+		optimistic: true,
+		ecmaVersion,
+		sourceType}).acquireAll(rootNode)[0];
+	const allScopes = {};
+	const stack = [globalScope];
+	while (stack.length) {
+		let scope = stack.pop();
+		const scopeId = scope.block.nodeId;
+		scope.block.isScopeBlock = true;
+		if (!allScopes[scopeId]) {
+			allScopes[scopeId] = scope;
+			stack.push(...scope.childScopes);
+		}
+	}
+	rootNode.allScopes = allScopes;
+	return allScopes;
+}
+
+/**
+ * @param {ASTNode} node
+ * @param {ASTScope[]} allScopes
+ * @return {ASTScope}
+ */
+function matchScopeToNode(node, allScopes) {
+	if (node.lineage?.length) {
+		for (const nid of [...node.lineage].reverse()) {
+			if (allScopes[nid]) {
+				let scope = allScopes[nid];
+				if (scope.type.includes('-name') && scope?.childScopes?.length === 1) scope = scope.childScopes[0];
+				return scope;
+			}
+		}
+	}
+	return allScopes[0]; // Global scope - this should never be reached
+}
+
+/**
+ *
+ * @param {string} inputCode
+ * @param {object} opts
+ * @return {Promise<ASTNode[]>}
+ */
+async function generateFlatASTAsync(inputCode, opts = {}) {
+	opts = { ...generateFlatASTDefaultOptions, ...opts };
+	const rootNode = generateRootNode(inputCode, opts);
+	const tree = extractNodesFromRoot(rootNode, opts);
+	const promises = [];
+	if (opts.detailed) {
+		const scopes = getAllScopes(rootNode);
+		for (let i = 0; i < tree.length; i++) {
+			promises.push(injectScopeToNodeAsync(tree[i], scopes));
+		}
+
 	}
+	return Promise.all(promises).then(() => tree);
 }
 
 module.exports = {
 	estraverse,
 	extractNodesFromRoot,
 	generateCode,
 	generateFlatAST,
+	generateFlatASTAsync,
 	generateRootNode,
+	injectScopeToNode,
+	injectScopeToNodeAsync,
 	parseCode,
 };

src/types.js

@@ -9,6 +9,7 @@ const {Scope} = require('eslint-scope');
  * @property {boolean} [async]
  * @property {ASTNode|ASTNode[]} [body]
  * @property {ASTNode} [callee]
+ * @property {ASTNode[]} [cases]
  * @property {ASTNode[]} [childNodes]
  * @property {boolean} [computed]
  * @property {ASTNode} [consequent]
@@ -17,6 +18,7 @@ const {Scope} = require('eslint-scope');
  * @property {ASTNode[]} [declarations]
  * @property {ASTNode} [declNode]
  * @property {boolean} [delegate]
+ * @property {ASTNode} [discriminant]
  * @property {ASTNode[]} [elements]
  * @property {number} [end]
  * @property {ASTNode} [exported]
@@ -28,6 +30,7 @@ const {Scope} = require('eslint-scope');
  * @property {ASTNode} [imported]
  * @property {ASTNode} [init]
  * @property {boolean} [isMarked]
+ * @property {boolean} [isScopeBlock]
  * @property {ASTNode} [key]
  * @property {string} [kind]
  * @property {ASTNode} [label]
@@ -75,6 +78,7 @@ class ASTNode {}
  * @property {ASTNode} block
  * @property {ASTScope[]} childScopes
  * @property {number} scopeId
+ * @property {string} type
  */
 class ASTScope extends Scope {}