From 30aa363c0c67d5ab2ced7c3033516067e5e282f3 Mon Sep 17 00:00:00 2001
From: Ben Baryo <ben.baryo@humansecurity.com>
Date: Tue, 12 Nov 2024 16:26:39 +0200
Subject: [PATCH] Add the typemap feature

---
 README.md             |  1 +
 src/flast.js          |  6 +++++-
 tests/parsing.test.js | 18 ++++++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 45fc263..0ae2c34 100644
--- a/README.md
+++ b/README.md
@@ -41,6 +41,7 @@ npm install
 - Tracks scope and connects each declaration to its references.  
   See [eslint-scope](https://github.com/eslint/eslint-scope) for more info on the scopes used.
 - Adds a unique id to each node to simplify tracking and understanding relations between nodes.
+- Maps the types to the node ids for easier access.
 - <u>Arborist</u> - marks nodes for replacement or deletion and applies all changes in a single iteration over the tree.
 
 ### flAST Data Structure
diff --git a/src/flast.js b/src/flast.js
index bd43896..01ea6ad 100644
--- a/src/flast.js
+++ b/src/flast.js
@@ -19,7 +19,7 @@ function parseCode(inputCode, opts = {}) {
 
 const excludedParentKeys = [
 	'type', 'start', 'end', 'range', 'sourceType', 'comments', 'srcClosure', 'nodeId',
-	'childNodes', 'parentNode', 'parentKey', 'scope',
+	'childNodes', 'parentNode', 'parentKey', 'scope', 'typeMap', 'lineage', 'allScopes',
 ];
 
 /**
@@ -123,6 +123,7 @@ function extractNodesFromRoot(rootNode, opts) {
 	opts = { ...generateFlatASTDefaultOptions, ...opts };
 	const tree = [];
 	let nodeId = 0;
+	const typeMap = {};
 
 	// noinspection JSUnusedGlobalSymbols
 	estraverse.traverse(rootNode, {
@@ -133,6 +134,8 @@ function extractNodesFromRoot(rootNode, opts) {
 		enter(node, parentNode) {
 			tree.push(node);
 			node.nodeId = nodeId++;
+			if (!typeMap[node.type]) typeMap[node.type] = [node.nodeId];
+			else typeMap[node.type].push(node.nodeId);
 			node.childNodes = [];
 			node.parentNode = parentNode;
 			node.parentKey = parentNode ? getParentKey(node) : '';
@@ -146,6 +149,7 @@ function extractNodesFromRoot(rootNode, opts) {
 			});
 		}
 	});
+	if (tree?.length) tree[0].typeMap = typeMap;
 	return tree;
 }
 
diff --git a/tests/parsing.test.js b/tests/parsing.test.js
index 9b21754..05fe5d0 100644
--- a/tests/parsing.test.js
+++ b/tests/parsing.test.js
@@ -55,4 +55,22 @@ describe('Parsing tests', () => {
 		const result = generateCode(ast[0]);
 		assert.strictEqual(result, expected);
 	});
+	it(`Verify the type map is generated accurately`, () => {
+		const code = `class a {
+  static b = 1;
+  #c = 2;
+}`;
+		const expected = {
+			Program: [0],
+			ClassDeclaration: [1],
+			Identifier: [2, 5],
+			ClassBody: [3],
+			PropertyDefinition: [4, 7],
+			Literal: [6, 9],
+			PrivateIdentifier: [8],
+		};
+		const ast = generateFlatAST(code);
+		const result = ast[0].typeMap;
+		assert.deepEqual(result, expected);
+	});
 });
\ No newline at end of file