From 9d2e9d4b03a9c63b930fa5581d407e4a909c2bd7 Mon Sep 17 00:00:00 2001
From: Ben Baryo <ben.baryo@humansecurity.com>
Date: Tue, 12 Nov 2024 16:32:01 +0200
Subject: [PATCH] Replace nodeIds with actual nodes

---
 README.md             |  2 +-
 src/flast.js          |  4 ++--
 src/types.js          |  1 +
 tests/parsing.test.js | 16 ++++++++--------
 4 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 0ae2c34..b846bc6 100644
--- a/README.md
+++ b/README.md
@@ -41,7 +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.
+- Maps the types to the nodes 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 01ea6ad..98a92b9 100644
--- a/src/flast.js
+++ b/src/flast.js
@@ -134,8 +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);
+			if (!typeMap[node.type]) typeMap[node.type] = [node];
+			else typeMap[node.type].push(node);
 			node.childNodes = [];
 			node.parentNode = parentNode;
 			node.parentKey = parentNode ? getParentKey(node) : '';
diff --git a/src/types.js b/src/types.js
index 7836095..b4f679f 100644
--- a/src/types.js
+++ b/src/types.js
@@ -73,6 +73,7 @@ import {Scope} from 'eslint-scope';
  * @property {ASTNode} [test]
  * @property {ASTNode} [tokens]
  * @property {Object[]} [trailingComments]
+ * @property {Object} [typeMap]
  * @property {ASTNode} [update]
  * @property {ASTNode|string|number|boolean} [value]
  */
diff --git a/tests/parsing.test.js b/tests/parsing.test.js
index 05fe5d0..962c2a6 100644
--- a/tests/parsing.test.js
+++ b/tests/parsing.test.js
@@ -60,16 +60,16 @@ describe('Parsing tests', () => {
   static b = 1;
   #c = 2;
 }`;
+		const ast = generateFlatAST(code);
 		const expected = {
-			Program: [0],
-			ClassDeclaration: [1],
-			Identifier: [2, 5],
-			ClassBody: [3],
-			PropertyDefinition: [4, 7],
-			Literal: [6, 9],
-			PrivateIdentifier: [8],
+			Program: [ast[0]],
+			ClassDeclaration: [ast[1]],
+			Identifier: [ast[2], ast[5]],
+			ClassBody: [ast[3]],
+			PropertyDefinition: [ast[4], ast[7]],
+			Literal: [ast[6], ast[9]],
+			PrivateIdentifier: [ast[8]],
 		};
-		const ast = generateFlatAST(code);
 		const result = ast[0].typeMap;
 		assert.deepEqual(result, expected);
 	});