From 9bfcad47b695b6970a34508a95362d5725edbc4e Mon Sep 17 00:00:00 2001
From: Ben Baryo <19845603+ctrl-escp@users.noreply.github.com>
Date: Fri, 5 Jul 2024 19:42:43 +0300
Subject: [PATCH] Utils (#26)

* Add utils:
  - applyIteratively
  - treeModifier
  - logger

* Add missing allScopes object

* Add utils tests

* Change imports to relative without __dirname
---
 src/arborist.js               |  2 +-
 src/flast.js                  |  1 +
 src/index.js                  |  7 ++--
 src/types.js                  |  1 +
 src/utils/applyIteratively.js | 65 +++++++++++++++++++++++++++++++++++
 src/utils/index.js            |  5 +++
 src/utils/logger.js           | 44 ++++++++++++++++++++++++
 src/utils/treeModifier.js     | 23 +++++++++++++
 tests/functionalityTests.js   |  1 +
 tests/parsingTests.js         |  2 +-
 tests/tester.js               |  7 ++--
 tests/utilsTests.js           | 32 +++++++++++++++++
 12 files changed, 182 insertions(+), 8 deletions(-)
 create mode 100644 src/utils/applyIteratively.js
 create mode 100644 src/utils/index.js
 create mode 100644 src/utils/logger.js
 create mode 100644 src/utils/treeModifier.js
 create mode 100644 tests/utilsTests.js

diff --git a/src/arborist.js b/src/arborist.js
index 4e5fd16..ff0e808 100644
--- a/src/arborist.js
+++ b/src/arborist.js
@@ -1,4 +1,4 @@
-const {generateCode, generateFlatAST,} = require(__dirname + '/flast');
+const {generateCode, generateFlatAST,} = require('./flast');
 
 const Arborist = class {
 	/**
diff --git a/src/flast.js b/src/flast.js
index 221d179..c6ccc05 100644
--- a/src/flast.js
+++ b/src/flast.js
@@ -79,6 +79,7 @@ function generateFlatAST(inputCode, opts = {}) {
 	if (opts.detailed) {
 		const scopes = getAllScopes(rootNode);
 		for (let i = 0; i < tree.length; i++) injectScopeToNode(tree[i], scopes);
+		tree[0].allScopes = scopes;
 	}
 	return tree;
 }
diff --git a/src/index.js b/src/index.js
index 8e3d0d8..f9f17c5 100644
--- a/src/index.js
+++ b/src/index.js
@@ -1,5 +1,6 @@
 module.exports = {
-	...require(__dirname + '/flast'),
-	...require(__dirname + '/arborist'),
-	...require(__dirname + '/types'),
+	...require('./flast'),
+	...require('./arborist'),
+	...require('./types'),
+	utils: require('./utils'),
 };
\ No newline at end of file
diff --git a/src/types.js b/src/types.js
index 07da911..d1b3fa8 100644
--- a/src/types.js
+++ b/src/types.js
@@ -3,6 +3,7 @@ const {Scope} = require('eslint-scope');
 /**
  * @typedef ASTNode
  * @property {string} type
+ * @property {object} [allScopes]
  * @property {ASTNode} [alternate]
  * @property {ASTNode} [argument]
  * @property {ASTNode[]} [arguments]
diff --git a/src/utils/applyIteratively.js b/src/utils/applyIteratively.js
new file mode 100644
index 0000000..e8ae15e
--- /dev/null
+++ b/src/utils/applyIteratively.js
@@ -0,0 +1,65 @@
+const {Arborist} = require('../arborist');
+const logger = require('./logger');
+const {createHash} = require('node:crypto');
+
+const generateHash = str => createHash('sha256').update(str).digest('hex');
+
+
+/**
+ * Apply functions to modify the script repeatedly until they are no long effective or the max number of iterations is reached.
+ * @param {string} script The target script to run the functions on.
+ * @param {function[]} funcs
+ * @param {number?} maxIterations (optional) Stop the loop after this many iterations at most.
+ * @return {string} The possibly modified script.
+ */
+function runLoop(script, funcs, maxIterations = 500) {
+	let scriptSnapshot = '';
+	let currentIteration = 0;
+	let changesCounter = 0;
+	let iterationsCounter = 0;
+	try {
+		let scriptHash = generateHash(script);
+		let arborist = new Arborist(script);
+		while (arborist.ast?.length && scriptSnapshot !== script && currentIteration < maxIterations) {
+			const iterationStartTime = Date.now();
+			scriptSnapshot = script;
+			// Mark each node with the script hash to distinguish cache of different scripts.
+			for (let i = 0; i < arborist.ast.length; i++) arborist.ast[i].scriptHash = scriptHash;
+			for (let i = 0; i <  funcs.length; i++) {
+				const func = funcs[i];
+				const funcStartTime = +new Date();
+				try {
+					logger.debug(`\t[!] Running ${func.name}...`);
+					arborist = func(arborist);
+					if (!arborist.ast?.length) break;
+					// If the hash doesn't exist it means the Arborist was replaced
+					const numberOfNewChanges = arborist.getNumberOfChanges() + +!arborist.ast[0].scriptHash;
+					if (numberOfNewChanges) {
+						changesCounter += numberOfNewChanges;
+						logger.log(`\t[+] ${func.name} applying ${numberOfNewChanges} new changes!`);
+						arborist.applyChanges();
+						script = arborist.script;
+						scriptHash = generateHash(script);
+						for (let j = 0; j < arborist.ast.length; j++) arborist.ast[j].scriptHash = scriptHash;
+					}
+				} catch (e) {
+					logger.error(`[-] Error in ${func.name} (iteration #${iterationsCounter}): ${e}\n${e.stack}`);
+				} finally {
+					logger.debug(`\t\t[!] Running ${func.name} completed in ` +
+              `${((+new Date() - funcStartTime) / 1000).toFixed(3)} seconds`);
+				}
+			}
+			++currentIteration;
+			++iterationsCounter;
+			logger.log(`[+] ==> Iteartion #${iterationsCounter} completed in ${(Date.now() - iterationStartTime) / 1000} seconds` +
+          ` with ${changesCounter ? changesCounter : 'no'} changes (${arborist.ast?.length || '???'} nodes)`);
+			changesCounter =  0;
+		}
+		if (changesCounter) script = arborist.script;
+	} catch (e) {
+		logger.error(`[-] Error on iteration #${iterationsCounter}: ${e}\n${e.stack}`);
+	}
+	return script;
+}
+
+module.exports = runLoop;
\ No newline at end of file
diff --git a/src/utils/index.js b/src/utils/index.js
new file mode 100644
index 0000000..7b85392
--- /dev/null
+++ b/src/utils/index.js
@@ -0,0 +1,5 @@
+module.exports = {
+	applyIteratively: require('./applyIteratively'),
+	logger: require('./logger'),
+	treeModifier: require('./treeModifier'),
+};
\ No newline at end of file
diff --git a/src/utils/logger.js b/src/utils/logger.js
new file mode 100644
index 0000000..545961b
--- /dev/null
+++ b/src/utils/logger.js
@@ -0,0 +1,44 @@
+const logLevels = {
+	DEBUG: 1,
+	LOG: 2,
+	ERROR: 3,
+	NONE: 9e10,
+};
+
+/**
+ * @param {number} logLevel
+ * @returns {function(*): void|undefined}
+ */
+function createLoggerForLevel(logLevel) {
+	if (!Object.values(logLevels).includes(logLevel)) throw new Error(`Unknown log level ${logLevel}.`);
+	return msg => logLevel >= logger.currentLogLevel ? logger.logFunc(msg) : undefined;
+}
+
+const logger = {
+	logLevels,
+	logFunc: console.log,
+	debug: createLoggerForLevel(logLevels.DEBUG),
+	log: createLoggerForLevel(logLevels.LOG),
+	error: createLoggerForLevel(logLevels.ERROR),
+	currentLogLevel: logLevels.NONE,
+
+	/**
+  * Set the current log level
+  * @param {number} newLogLevel
+  */
+	setLogLevel(newLogLevel) {
+		if (!Object.values(this.logLevels).includes(newLogLevel)) throw new Error(`Unknown log level ${newLogLevel}.`);
+		this.currentLogLevel = newLogLevel;
+	},
+
+	setLogLeveNone() {this.setLogLevel(this.logLevels.NONE);},
+	setLogLeveDebug() {this.setLogLevel(this.logLevels.DEBUG);},
+	setLogLeveLog() {this.setLogLevel(this.logLevels.LOG);},
+	setLogLeveError() {this.setLogLevel(this.logLevels.ERROR);},
+  
+	setLogFunc(newLogfunc) {
+		this.logFunc = newLogfunc;
+	},
+};
+
+module.exports = logger;
\ No newline at end of file
diff --git a/src/utils/treeModifier.js b/src/utils/treeModifier.js
new file mode 100644
index 0000000..4a6cc94
--- /dev/null
+++ b/src/utils/treeModifier.js
@@ -0,0 +1,23 @@
+/**
+ * Boilerplate for filter functions that identify the desired structure and a modifier function that modifies the tree.
+ * An optional name for the function can be provided for better logging.
+ * @param {Function} filterFunc
+ * @param {Function} modFunc
+ * @param {string} [funcName]
+ * @returns {function(Arborist): Arborist}
+ */
+function treeModifier(filterFunc, modFunc, funcName) {
+	const func = function(arb) {
+		for (let i = 0; i < arb.ast.length; i++) {
+			const n = arb.ast[i];
+			if (filterFunc(n, arb)) {
+				modFunc(n, arb);
+			}
+		}
+		return arb;
+	};
+	if (funcName) Object.defineProperty(func, 'name', {value: funcName});
+	return func;
+}
+
+module.exports = treeModifier;
\ No newline at end of file
diff --git a/tests/functionalityTests.js b/tests/functionalityTests.js
index 3a5819d..b0014ad 100644
--- a/tests/functionalityTests.js
+++ b/tests/functionalityTests.js
@@ -38,6 +38,7 @@ module.exports = [
 				'generateCode',
 				'generateFlatAST',
 				'parseCode',
+				'utils',
 			];
 			function tryImporting(importName) {
 				const {[importName]: tempImport} = require(importSource);
diff --git a/tests/parsingTests.js b/tests/parsingTests.js
index d102804..a46e3d6 100644
--- a/tests/parsingTests.js
+++ b/tests/parsingTests.js
@@ -1,5 +1,5 @@
 const assert = require('node:assert');
-const {generateFlatAST} = require(__dirname + '/../src/index');
+const {generateFlatAST} = require('../src/index');
 
 module.exports = [
 	{
diff --git a/tests/tester.js b/tests/tester.js
index ecb273b..f073246 100644
--- a/tests/tester.js
+++ b/tests/tester.js
@@ -1,7 +1,8 @@
 const tests = {
-	Parsing: __dirname + '/parsingTests',
-	Functionality: __dirname + '/functionalityTests',
-	Arborist: __dirname + '/aboristTests',
+	Parsing: './parsingTests',
+	Functionality: './functionalityTests',
+	Arborist: './aboristTests',
+	Utils: './utilsTests',
 };
 
 let allTests = 0;
diff --git a/tests/utilsTests.js b/tests/utilsTests.js
new file mode 100644
index 0000000..d3d5692
--- /dev/null
+++ b/tests/utilsTests.js
@@ -0,0 +1,32 @@
+const {utils} = require(__dirname + '/../src/index');
+const assert = require('node:assert');
+module.exports = [
+	{
+		enabled: true,
+		name: 'treeModifier + applyIteratively',
+		description: '',
+		run() {
+			const code = `console.log('Hello' + ' ' + 'there');`;
+			const expectedOutput = `console.log('General' + ' ' + 'Kenobi');`;
+			const expectedFuncName = 'StarWarsDialog';
+			const replacements = {
+				Hello: 'General',
+				there: 'Kenobi',
+			};
+			let result = code;
+			const f = n => n.type === 'Literal' && replacements[n.value];
+			const m = (n, arb) => arb.markNode(n, {
+				type: 'Literal',
+				value: replacements[n.value],
+			});
+			const generatedFunc = utils.treeModifier(f, m, expectedFuncName);
+			result = utils.applyIteratively(result, [generatedFunc]);
+
+			assert.equal(result, expectedOutput,
+				`Result does not match expected output.`);
+			assert.equal(generatedFunc.name, expectedFuncName,
+				`The name of the generated function does not match.`);
+			return true;
+		},
+	},
+];
\ No newline at end of file