diff --git a/.gitignore b/.gitignore index d0ff856..20ac21b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,7 @@ +.vscode/ +.cursor node_modules/ .idea/ *tmp*/ -*tmp*.* \ No newline at end of file +*tmp*.* +.DS_Store diff --git a/README.md b/README.md index b846bc6..72cb478 100644 --- a/README.md +++ b/README.md @@ -2,220 +2,242 @@ [![Run Tests](https://github.com/PerimeterX/flast/actions/workflows/node.js.yml/badge.svg?branch=main)](https://github.com/PerimeterX/flast/actions/workflows/node.js.yml) [![Downloads](https://img.shields.io/npm/dm/flast.svg?maxAge=43200)](https://www.npmjs.com/package/flast) -Flatten an Abstract Syntax Tree by placing all the nodes in a single flat array. +Efficient, flat, and richly-annotated JavaScript AST manipulation for code transformation, analysis, and more. -For comments and suggestions feel free to open an issue or find me on Twitter - [@ctrl__esc](https://twitter.com/ctrl__esc) +For comments and suggestions feel free to open an issue or find me on Twitter/X - [@ctrl__esc](https://x.com/ctrl__esc) ## Table of Contents * [Installation](#installation) - * [npm](#npm) - * [Clone The Repo](#clone-the-repo) * [Features](#features) - * [flAST Data Structure](#flast-data-structure) -* [Usage](#usage) - * [flAST](#flast) - * [generateFlatAST Options](#generateflatast-options) - * [generateCode Options](#generatecode-options) - * [Arborist](#arborist) +* [Usage Examples](#usage-examples) * [How to Contribute](#how-to-contribute) -*** +* [Projects Using flAST](#projects-using-flast) + +--- ## Installation +Requires Node 18 or newer. + ### npm ```bash npm install flast ``` -### Clone The Repo -Requires Node 18 or newer. +### Clone the Repo ```bash git clone git@github.com:PerimeterX/flast.git cd flast npm install ``` - -*** +--- ## Features -- Keeps all relations between parent and child nodes. -- 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 nodes for easier access. -- Arborist - marks nodes for replacement or deletion and applies all changes in a single iteration over the tree. - -### flAST Data Structure -
- Example of how a flat AST would look like. - -Input code: `console.log('flAST');`. -Output object: -```javascript -const tree = [ - { - type: 'program', - start: 0, - end: 21, - range: [0, 21], - body: [ - '' - ], - sourceType: 'script', - comments: [], - nodeId: 0, - src: "console.log('flAST');", - childNodes: [ - '' - ], - parentNode: null, - scope: '' - }, - { - type: 'ExpressionStatement', - start: 0, - end: 21, - range: [0, 21], - expression: '', - nodeId: 1, - src: "console.log('flAST');", - childNodes: [ - '' - ], - parentNode: '', - scope: '' - }, - { - type: 'CallExpression', - start: 0, - end: 20, - range: [0, 20], - callee: '', - arguments: [ - '' - ], - optional: false, - nodeId: 2, - src: "console.log('flAST')", - childNodes: [ - '', - '' - ], - parentNode: '', - scope: '' - }, - { - type: 'MemberExpression', - start: 0, - end: 11, - range: [0, 11], - object: '', - property: '', - computed: false, - optional: false, - nodeId: 3, - src: 'console.log', - childNodes: [ - '', - '' - ], - parentNode: '', - scope: '' - }, - { - type: 'Identifier', - start: 0, - end: 7, - range: [0, 7], - name: 'console', - nodeId: 4, - src: 'console', - childNodes: [], - parentNode: '', - scope: '' - }, - { - type: 'Identifier', - start: 8, - end: 11, - range: [8, 11], - name: 'log', - nodeId: 5, - src: 'log', - childNodes: [], - parentNode: '', - scope: '' - }, - { - type: 'Literal', - start: 12, - end: 19, - range: [12, 19], - value: "flAST", - raw: "'flAST'", - nodeId: 6, - src: "'flAST'", - childNodes: [], - parentNode: '', - scope: '' + +### Parsing and Code Generation +- **Code to AST:** Parse JavaScript code into a flat, richly annotated AST. +- **AST to Code:** Generate code from any AST node, supporting round-trip transformations. + +### Flat AST Structure +- **Flat AST (`generateFlatAST`):** All nodes are in a single array, allowing direct access and efficient traversal without recursive tree-walking. +- **Type Map:** `ast[0].typeMap` provides fast lookup of all nodes by type. +- **Scopes:** `ast[0].allScopes` gives direct access to all lexical scopes. + +### Node Richness +Each node in the flat AST includes: +- `src`: The original code for this node. +- `parentNode` and `childNodes`: Easy navigation and context. +- `parentKey`: The property name this node occupies in its parent. +- `declNode`: For variables, a reference to their declaration node. +- `references`: For declarations, a list of all reference nodes. +- `lineage`: Traceable ancestry of scopes for each node. +- `nodeId`: Unique identifier for each node. +- `scope`, `scopeId`, and more for advanced analysis. + +### Arborist: AST Modification +- **Delete nodes:** Mark nodes for removal and apply changes safely. +- **Replace nodes:** Mark nodes for replacement, with all changes validated and applied in a single pass. + +### Utilities +- **applyIteratively:** Apply a series of transformation functions (using Arborist) to the AST/code, iterating until no further changes are made. Automatically reverts changes that break the code. +- **logger:** Simple log utility that can be controlled downstream and used for debugging or custom output. +- **treeModifier:** (Deprecated) Simple wrapper for AST iteration. +--- + +## Usage Examples + +> **Tip:** +> For best performance, always iterate over only the relevant node types using `ast[0].typeMap`. For example, to process all identifiers and variable declarations: +> ```js +> const relevantNodes = [ +> ...ast[0].typeMap.Identifier, +> ...ast[0].typeMap.VariableDeclarator, +> ]; +> for (let i = 0; i < relevantNodes.length; i++) { +> const n = relevantNodes[i]; +> // ... process n ... +> } +> ``` +> Only iterate over the entire AST as a last resort. + +### Basic Example +```js +import {Arborist} from 'flast'; + +const replacements = {'Hello': 'General', 'there!': 'Kenobi'}; + +const arb = new Arborist(`console.log('Hello' + ' ' + 'there!');`); +// This is equivalent to: +// const ast = generateFlatAST(`console.log('Hello' + ' ' + 'there!');`); +// const arb = new Arborist(ast); +// Since the Arborist accepts either code as a string or a flat AST object. + +for (let i = 0; i < arb.ast.length; i++) { + const n = arb.ast[i]; + if (n.type === 'Literal' && replacements[n.value]) { + arb.markNode(n, { + type: 'Literal', + value: replacements[n.value], + raw: `'${replacements[n.value]}'`, + }); } -]; +} +arb.applyChanges(); +console.log(arb.script); // console.log('General' + ' ' + 'Kenobi'); ``` -
+--- + +### Example 1: Numeric Calculation Simplification +Replace constant numeric expressions with their computed value. +```js +import {applyIteratively} from 'flast'; -## Usage -### flAST +function simplifyNumericExpressions(arb) { + const binaryNodes = arb.ast[0].typeMap.BinaryExpression || []; + for (let i = 0; i < binaryNodes.length; i++) { + const n = binaryNodes[i]; + if (n.left.type === 'Literal' && typeof n.left.value === 'number' && + n.right.type === 'Literal' && typeof n.right.value === 'number') { + let result; + switch (n.operator) { + case '+': result = n.left.value + n.right.value; break; + case '-': result = n.left.value - n.right.value; break; + case '*': result = n.left.value * n.right.value; break; + case '/': result = n.left.value / n.right.value; break; + default: continue; + } + arb.markNode(n, {type: 'Literal', value: result, raw: String(result)}); + } + } + return arb; +} -```javascript -import {generateFlatAST, generateCode} from 'flast'; -const ast = generateFlatAST(`console.log('flAST')`); -const reconstructedCode = generateCode(ast[0]); // rebuild from root node +const script = 'let x = 5 * 3 + 1;'; +const result = applyIteratively(script, [simplifyNumericExpressions]); +console.log(result); // let x = 16; ``` -#### generateFlatAST Options -```javascript -const generateFlatASTDefaultOptions = { - detailed: true, // If false, include only original node without any further details - includeSrc: true, // If false, do not include node src. Only available when `detailed` option is true +--- + +### Example 2: Transform Arrow Function to Regular Function +```js +import {applyIteratively} from 'flast'; + +function arrowToFunction(arb) { + const arrowNodes = arb.ast[0].typeMap.ArrowFunctionExpression || []; + for (let i = 0; i < arrowNodes.length; i++) { + const n = arrowNodes[i]; + arb.markNode(n, { + type: 'FunctionExpression', + id: null, + params: n.params, + body: n.body.type === 'BlockStatement' ? n.body : {type: 'BlockStatement', body: [{ type: 'ReturnStatement', argument: n.body }] }, + generator: false, + async: n.async, + expression: false, + }); + } + return arb; +} + +const script = 'const f = (a, b) => a + b;'; +const result = applyIteratively(script, [arrowToFunction]); +console.log(result); +/* +const f = function(a, b) { + return a + b; }; +*/ ``` +--- -#### generateCode Options -See [Espree's documentation](https://github.com/eslint/espree#options) for more information -```javascript -const generateCodeDefaultOptions = { - format: { - indent: { - style: ' ', - adjustMultilineComment: true, - }, - quotes: 'auto', - escapeless: true, - compact: false, - }, - comment: true, -}; +### Example 3: Modify Nodes Based on Attached Comments +Suppose you want to double any numeric literal that has a comment `// double` attached: +```js +import {applyIteratively} from 'flast'; + +function doubleLiteralsWithComment(arb) { + const literalNodes = arb.ast[0].typeMap.Literal || []; + for (let i = 0; i < literalNodes.length; i++) { + const n = literalNodes[i]; + if ( + typeof n.value === 'number' && + n.leadingComments && + n.leadingComments.some(c => c.value.includes('double')) + ) { + arb.markNode(n, { type: 'Literal', value: n.value * 2, raw: String(n.value * 2) }); + } + } + return arb; +} + +const script = 'const x = /* double */ 21;'; +const result = applyIteratively(script, [doubleLiteralsWithComment], 1); // Last argument is the maximum number of iterations allowed. +console.log(result); // const x = /* double */ 42; ``` +--- -### Arborist +### Example 4: Proxy Variable Replacement Using References +Replace all references to a variable that is a proxy for another variable. +```js +import {applyIteratively} from 'flast'; -```javascript -import {generateFlatAST, generateCode, Arborist} from 'flast'; -const ast = generateFlatAST(`console.log('Hello' + ' ' + 'there!');`); -const replacements = { - 'Hello': 'General', - 'there!': 'Kenobi', -}; -const arborist = new Arborist(ast); -// Mark all relevant nodes for replacement. -ast.filter(n => n.type === 'Literal' && replacements[n.value]).forEach(n => arborist.markNode(n, { - type: 'Literal', - value: replacements[n.value], - raw: `'${replacements[n.value]}'`, -})); -const numberOfChangesMade = arborist.applyChanges(); -console.log(generateCode(arborist.ast[0])); // console.log('General' + ' ' + 'Kenobi'); +function replaceProxyVars(arb) { + const declarators = arb.ast[0].typeMap.VariableDeclarator || []; + for (let i = 0; i < declarators.length; i++) { + const n = declarators[i]; + if (n.init && n.init.type === 'Identifier' && n.id && n.id.name) { + // Replace all references to this variable with the variable it proxies + const refs = n.references || []; + for (let j = 0; j < refs.length; j++) { + const ref = refs[j]; + arb.markNode(ref, { + type: 'Identifier', + name: n.init.name, + }); + } + } + } + return arb; +} + +const script = 'var a = b; var b = 42; console.log(a);'; +const result = applyIteratively(script, [replaceProxyVars]); +console.log(result); // var a = b; var b = 42; console.log(b); ``` -The Arborist can be called with an extra argument - logFunc - which can be used to override the log -function inside the arborist. +--- + +## Projects Using flAST +- **[Obfuscation-Detector](https://github.com/PerimeterX/obfuscation-detector):** Uses flAST to analyze and detect unique obfuscation in JavaScript code. +- **[REstringer](https://github.com/PerimeterX/restringer):** Uses flAST for advanced code transformation and analysis. + +--- + +## Best Practices + +### AST Mutability +You can directly mutate nodes in the flat AST (e.g., changing properties, adding or removing nodes). However, for safety and script validity, it's best to use the Arborist for all structural changes. The Arborist verifies your changes and prevents breaking the code, ensuring that the resulting AST remains valid and that all node information is updated correctly. + +- **Direct mutation is possible**, but should be used with caution. +- **Recommended:** Use the Arborist's `markNode` method for all node deletions and replacements. ## How to Contribute -To contribute to this project see our [contribution guide](CONTRIBUTING.md) \ No newline at end of file +To contribute to this project see our [contribution guide](CONTRIBUTING.md) diff --git a/src/arborist.js b/src/arborist.js index 52c5bde..dee7165 100644 --- a/src/arborist.js +++ b/src/arborist.js @@ -1,9 +1,13 @@ import {logger} from './utils/logger.js'; import {generateCode, generateFlatAST} from './flast.js'; -const Arborist = class { +/** + * Arborist allows marking nodes for deletion or replacement, and then applying all changes in a single pass. + * Note: Marking a node with markNode() only sets a flag; the AST is not officially changed until applyChanges() is called. + */ +class Arborist { /** - * @param {string|ASTNode[]} scriptOrFlatAstArr - the target script or a flat AST array + * @param {string|ASTNode[]} scriptOrFlatAstArr - The target script or a flat AST array. */ constructor(scriptOrFlatAstArr) { this.script = ''; @@ -46,10 +50,9 @@ const Arborist = class { } /** - * Replace the target node with another node or delete the target node completely, depending on whether a replacement - * node is provided. + * Mark a node for replacement or deletion. This only sets a flag; the AST is not changed until applyChanges() is called. * @param {ASTNode} targetNode The node to replace or remove. - * @param {object|ASTNode} replacementNode If exists, replace the target node with this node. + * @param {object|ASTNode} [replacementNode] If exists, replace the target node with this node. */ markNode(targetNode, replacementNode) { if (!targetNode.isMarked) { @@ -67,28 +70,48 @@ const Arborist = class { } } + /** + * Merge comments from a source node into a target node or array. + * @param {ASTNode|Object} target - The node or array element to receive comments. + * @param {ASTNode} source - The node whose comments should be merged. + * @param {'leadingComments'|'trailingComments'} which + */ + static mergeComments(target, source, which) { + if (!source[which] || !source[which].length) return; + if (!target[which]) { + target[which] = [...source[which]]; + } else if (target[which] !== source[which]) { + target[which] = target[which].concat(source[which]); + } + } + /** * Iterate over the complete AST and replace / remove marked nodes, * then rebuild code and AST to validate changes. + * + * Note: If you delete a node that is the only child of its parent (e.g., the only statement in a block), + * you may leave the parent in an invalid or empty state. Consider cleaning up empty parents if needed. + * * @return {number} The number of modifications made. */ applyChanges() { let changesCounter = 0; + let rootNode = this.ast[0]; try { if (this.getNumberOfChanges() > 0) { - let rootNode = this.ast[0]; if (rootNode.isMarked) { const rootNodeReplacement = this.replacements.find(n => n[0].nodeId === 0); ++changesCounter; this.logger.debug(`[+] Applying changes to the root node...`); - const leadingComments = rootNode.leadingComments || []; + const leadingComments = rootNode.leadingComments || []; const trailingComments = rootNode.trailingComments || []; rootNode = rootNodeReplacement[1]; - if (leadingComments.length && rootNode.leadingComments !== leadingComments) rootNode.leadingComments = (rootNode.leadingComments || []).concat(leadingComments); - if (trailingComments.length && rootNode.trailingComments !== trailingComments) rootNode.trailingComments = (rootNode.trailingComments || []).concat(trailingComments); + if (leadingComments.length && rootNode.leadingComments !== leadingComments) + Arborist.mergeComments(rootNode, {leadingComments}, 'leadingComments'); + if (trailingComments.length && rootNode.trailingComments !== trailingComments) + Arborist.mergeComments(rootNode, {trailingComments}, 'trailingComments'); } else { - for (let i = 0; i < this.markedForDeletion.length; i++) { - const targetNodeId = this.markedForDeletion[i]; + for (const targetNodeId of this.markedForDeletion) { try { let targetNode = this.ast[targetNodeId]; targetNode = targetNode.nodeId === targetNodeId ? targetNode : this.ast.find(n => n.nodeId === targetNodeId); @@ -96,47 +119,54 @@ const Arborist = class { const parent = targetNode.parentNode; if (parent[targetNode.parentKey] === targetNode) { delete parent[targetNode.parentKey]; - const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []); - if (comments.length) parent.trailingComments = (parent.trailingComments || []).concat(comments); + Arborist.mergeComments(parent, targetNode, 'trailingComments'); ++changesCounter; } else if (Array.isArray(parent[targetNode.parentKey])) { const idx = parent[targetNode.parentKey].indexOf(targetNode); - parent[targetNode.parentKey].splice(idx, 1); - const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []); - if (comments.length) { - const targetParent = idx > 0 ? parent[targetNode.parentKey][idx - 1] : parent[targetNode.parentKey].length > 1 ? parent[targetNode.parentKey][idx + 1] : parent; - targetParent.trailingComments = (targetParent.trailingComments || []).concat(comments); + if (idx !== -1) { + parent[targetNode.parentKey].splice(idx, 1); + const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []); + let targetParent = null; + if (parent[targetNode.parentKey].length > 0) { + if (idx > 0) { + targetParent = parent[targetNode.parentKey][idx - 1]; + Arborist.mergeComments(targetParent, {trailingComments: comments}, 'trailingComments'); + } else { + targetParent = parent[targetNode.parentKey][0]; + Arborist.mergeComments(targetParent, {leadingComments: comments}, 'leadingComments'); + } + } else { + this.logger.debug(`[!] Deleted last element from array '${targetNode.parentKey}' in parent node type '${parent.type}'. Array is now empty.`); + Arborist.mergeComments(parent, {trailingComments: comments}, 'trailingComments'); + } + ++changesCounter; } - ++changesCounter; } } } catch (e) { this.logger.debug(`[-] Unable to delete node: ${e}`); } } - for (let i = 0; i < this.replacements.length; i++) { - const [targetNode, replacementNode] = this.replacements[i]; + for (const [targetNode, replacementNode] of this.replacements) { try { if (targetNode) { const parent = targetNode.parentNode; if (parent[targetNode.parentKey] === targetNode) { parent[targetNode.parentKey] = replacementNode; - const leadingComments = targetNode.leadingComments || []; - const trailingComments = targetNode.trailingComments || []; - if (leadingComments.length) replacementNode.leadingComments = (replacementNode.leadingComments || []).concat(leadingComments); - if (trailingComments.length) replacementNode.trailingComments = (replacementNode.trailingComments || []).concat(trailingComments); + Arborist.mergeComments(replacementNode, targetNode, 'leadingComments'); + Arborist.mergeComments(replacementNode, targetNode, 'trailingComments'); ++changesCounter; } else if (Array.isArray(parent[targetNode.parentKey])) { const idx = parent[targetNode.parentKey].indexOf(targetNode); parent[targetNode.parentKey][idx] = replacementNode; const comments = (targetNode.leadingComments || []).concat(targetNode.trailingComments || []); if (idx > 0) { - const commentsTarget = parent[targetNode.parentKey][idx - 1]; - commentsTarget.trailingComments = (commentsTarget.trailingComments || []).concat(comments); + Arborist.mergeComments(parent[targetNode.parentKey][idx - 1], {trailingComments: comments}, 'trailingComments'); } else if (parent[targetNode.parentKey].length > 1) { - const commentsTarget = parent[targetNode.parentKey][idx + 1]; - commentsTarget.leadingComments = (commentsTarget.leadingComments || []).concat(comments); - } else parent.trailingComments = (parent.trailingComments || []).concat(comments); + Arborist.mergeComments(parent[targetNode.parentKey][idx + 1], {leadingComments: comments}, 'leadingComments'); + } else { + Arborist.mergeComments(parent, {trailingComments: comments}, 'trailingComments'); + } ++changesCounter; } } @@ -145,21 +175,21 @@ const Arborist = class { } } } - if (changesCounter) { - this.replacements.length = 0; - this.markedForDeletion.length = 0; - // If any of the changes made will break the script the next line will fail and the - // script will remain the same. If it doesn't break, the changes are valid and the script can be marked as modified. - const script = generateCode(rootNode); - const ast = generateFlatAST(script); - if (ast && ast.length) { - this.ast = ast; - this.script = script; - } - else { - this.logger.log(`[-] Modified script is invalid. Reverting ${changesCounter} changes...`); - changesCounter = 0; - } + } + if (changesCounter) { + this.replacements.length = 0; + this.markedForDeletion.length = 0; + // If any of the changes made will break the script the next line will fail and the + // script will remain the same. If it doesn't break, the changes are valid and the script can be marked as modified. + const script = generateCode(rootNode); + const ast = generateFlatAST(script); + if (ast && ast.length) { + this.ast = ast; + this.script = script; + } + else { + this.logger.log(`[-] Modified script is invalid. Reverting ${changesCounter} changes...`); + changesCounter = 0; } } } catch (e) { @@ -168,7 +198,7 @@ const Arborist = class { ++this.appliedCounter; return changesCounter; } -}; +} export { Arborist, diff --git a/src/flast.js b/src/flast.js index 0b42c73..dc16df2 100644 --- a/src/flast.js +++ b/src/flast.js @@ -29,7 +29,7 @@ const generateFlatASTDefaultOptions = { // If false, do not include node src includeSrc: true, // Retry to parse the code with sourceType: 'script' if 'module' failed with 'strict' error message - alernateSourceTypeOnFailure: true, + alternateSourceTypeOnFailure: true, // Options for the espree parser parseOpts: { sourceType, @@ -78,6 +78,7 @@ const generateCodeDefaultOptions = { /** * @param {ASTNode} rootNode * @param {object} opts Optional changes to behavior. See generateCodeDefaultOptions for available options. +* All escodegen options are supported, including sourceMap, sourceMapWithCode, etc. * @return {string} Code generated from AST */ function generateCode(rootNode, opts = {}) { @@ -97,8 +98,17 @@ function generateRootNode(inputCode, opts = {}) { rootNode = parseCode(inputCode, parseOpts); if (opts.includeSrc) rootNode.srcClosure = createSrcClosure(inputCode); } catch (e) { - if (opts.alernateSourceTypeOnFailure && e.message.includes('in strict mode')) rootNode = parseCode(inputCode, {...parseOpts, sourceType: 'script'}); - else logger.debug(e); + // If any parse error occurs and alternateSourceTypeOnFailure is set, try 'script' mode + if (opts.alternateSourceTypeOnFailure) { + try { + rootNode = parseCode(inputCode, {...parseOpts, sourceType: 'script'}); + if (opts.includeSrc) rootNode.srcClosure = createSrcClosure(inputCode); + } catch (e2) { + logger.debug('Failed to parse as module and script:', e, e2); + } + } else { + logger.debug(e); + } } return rootNode; } @@ -168,41 +178,65 @@ function extractNodesFromRoot(rootNode, opts) { } if (opts.detailed) { const identifiers = typeMap.Identifier || []; + const scopeVarMaps = buildScopeVarMaps(scopes); for (let i = 0; i < identifiers.length; i++) { - mapIdentifierRelations(identifiers[i]); + mapIdentifierRelations(identifiers[i], scopeVarMaps); } } - if (allNodes?.length) allNodes[0].typeMap = typeMap; + if (allNodes?.length) { + allNodes[0].typeMap = new Proxy(typeMap, { + get(target, prop, receiver) { + if (prop in target) { + return Reflect.get(target, prop, receiver); + } + return []; // Return an empty array for any undefined type + }, + }); + } return allNodes; } +/** + * Precompute a map of variable names to declarations for each scope for fast lookup. + * @param {object} scopes + * @return {Map} Map of scopeId to { [name]: variable } + */ +function buildScopeVarMaps(scopes) { + const scopeVarMaps = {}; + for (const scopeId in scopes) { + const scope = scopes[scopeId]; + const varMap = {}; + for (let i = 0; i < scope.variables.length; i++) { + const v = scope.variables[i]; + varMap[v.name] = v; + } + scopeVarMaps[scopeId] = varMap; + } + return scopeVarMaps; +} + /** * @param {ASTNode} node + * @param {object} scopeVarMaps */ -function mapIdentifierRelations(node) { +function mapIdentifierRelations(node, scopeVarMaps) { // Track references and declarations // Prevent assigning declNode to member expression properties or object keys if (node.type === 'Identifier' && !(!node.parentNode.computed && ['property', 'key'].includes(node.parentKey))) { - const variables = []; - for (let i = 0; i < node.scope.variables.length; i++) { - if (node.scope.variables[i].name === node.name) variables.push(node.scope.variables[i]); - } - if (node.parentKey === 'id' || variables?.[0]?.identifiers?.includes(node)) { + const scope = node.scope; + const varMap = scope && scopeVarMaps ? scopeVarMaps[scope.scopeId] : undefined; + const variable = varMap ? varMap[node.name] : undefined; + if (node.parentKey === 'id' || variable?.identifiers?.includes(node)) { node.references = node.references || []; } else { // Find declaration by finding the closest declaration of the same name. let decls = []; - if (variables?.length) { - for (let i = 0; i < variables.length; i++) { - if (variables[i].name === node.name) { - decls = variables[i].identifiers || []; - break; - } - } - } else { - for (let i = 0; i < node.scope.references.length; i++) { - if (node.scope.references[i].identifier.name === node.name) { - decls = node.scope.references[i].resolved?.identifiers || []; + if (variable) { + decls = variable.identifiers || []; + } else if (scope && scope.references) { + for (let i = 0; i < scope.references.length; i++) { + if (scope.references[i].identifier.name === node.name) { + decls = scope.references[i].resolved?.identifiers || []; break; } } diff --git a/src/types.js b/src/types.js index cd5d9b9..0a205a3 100644 --- a/src/types.js +++ b/src/types.js @@ -1,91 +1,100 @@ import {Scope} from 'eslint-scope'; +/** + * @typedef {Object.} ASTTypeMap - Map of node type to array of nodes. + */ + +/** + * @typedef {Object.} ASTAllScopes - Map of scopeId to ASTScope. + */ + /** * @typedef ASTNode - * @property {string} type - * @property {object} [allScopes] - * @property {ASTNode} [alternate] - * @property {ASTNode} [argument] - * @property {ASTNode[]} [arguments] - * @property {boolean} [async] - * @property {ASTNode|ASTNode[]} [body] - * @property {ASTNode} [callee] - * @property {ASTNode[]} [cases] - * @property {ASTNode[]} [childNodes] - * @property {Object[]} [comments] - * @property {boolean} [computed] - * @property {ASTNode} [consequent] - * @property {string} [cooked] - * @property {ASTNode} [declaration] - * @property {ASTNode[]} [declarations] - * @property {ASTNode} [declNode] - * @property {boolean} [delegate] - * @property {ASTNode} [discriminant] - * @property {ASTNode[]} [elements] - * @property {number} [end] - * @property {ASTNode} [exported] - * @property {ASTNode|boolean} [expression] - * @property {ASTNode[]} [expressions] - * @property {string} [flags] - * @property {boolean} [generator] - * @property {ASTNode} [id] - * @property {ASTNode} [imported] - * @property {ASTNode} [init] - * @property {boolean} [isEmpty] True when the node is set for deletion but should be replced with an Empty Statement instead - * @property {boolean} [isMarked] True when the node has already been marked for replacement or deletion - * @property {ASTNode} [key] - * @property {string} [kind] - * @property {ASTNode} [label] - * @property {Object[]} [leadingComments] - * @property {ASTNode} [left] - * @property {number[]} [lineage] The nodeIds of all parent nodes - * @property {ASTNode} [local] - * @property {boolean} [method] - * @property {string} [name] - * @property {number} [nodeId] A unique id in the AST - * @property {string} [operator] - * @property {ASTNode} [object] - * @property {string} [pattern] - * @property {ASTNode[]} [params] - * @property {string} [parentKey] The designation the node has within the parent node - * @property {ASTNode} [parentNode] - * @property {boolean} [prefix] - * @property {ASTNode} [property] - * @property {ASTNode[]} [properties] - * @property {ASTNode[]} [quasis] - * @property {number[]} [range] - * @property {string} [raw] - * @property {ASTNode[]} [references] - * @property {ASTNode} [regex] - * @property {ASTNode} [right] - * @property {ASTScope} [scope] - * @property {number} [scopeId] For nodes which are also a scope's block - * @property {string} [scriptHash] - * @property {boolean} [shorthand] - * @property {ASTNode} [source] - * @property {string} [sourceType] - * @property {ASTNode[]} [specifiers] - * @property {boolean} [static] - * @property {number} [start] - * @property {string|function} [src] The source code for the node - * @property {ASTNode} [superClass] - * @property {boolean} [tail] - * @property {ASTNode} [test] - * @property {ASTNode} [tokens] - * @property {Object[]} [trailingComments] - * @property {Object} [typeMap] - * @property {ASTNode} [update] - * @property {ASTNode|string|number|boolean} [value] + * @property {ASTNode[]} childNodes - Array of child nodes. + * @property {number} nodeId - Unique id in the AST. + * @property {string} parentKey - The property name this node occupies in its parent. + * @property {ASTNode|null} parentNode - Parent node, or null for the root. + * @property {ASTTypeMap} typeMap - Only present on the root node. Map of node type to array of nodes. + * @property {string} type - Node type. + * + * @property {ASTAllScopes} [allScopes] - Only present on the root node. Map of scopeId to ASTScope. + * @property {ASTNode} [alternate] - Alternate branch (e.g., in IfStatement). + * @property {ASTNode} [argument] - Argument node. + * @property {ASTNode[]} [arguments] - Array of argument nodes. + * @property {boolean} [async] - True if the function is async. + * @property {ASTNode|ASTNode[]} [body] - Function or block body. + * @property {ASTNode} [callee] - Callee node in a CallExpression. + * @property {ASTNode[]} [cases] - Switch cases. + * @property {Object[]} [comments] - Comments attached to the node. + * @property {boolean} [computed] - True if property is computed. + * @property {ASTNode} [consequent] - Consequent branch (e.g., in IfStatement). + * @property {string} [cooked] - Cooked value for template literals. + * @property {ASTNode} [declaration] - Declaration node. + * @property {ASTNode[]} [declarations] - Array of declaration nodes. + * @property {ASTNode} [declNode] - Only present on identifier nodes that are references (detailed: true). + * @property {boolean} [delegate] - True if yield*. + * @property {ASTNode} [discriminant] - Switch discriminant. + * @property {ASTNode[]} [elements] - Array elements. + * @property {number} [end] - End position in source. + * @property {ASTNode} [exported] - Exported node. + * @property {ASTNode|boolean} [expression] - Expression node or boolean. + * @property {ASTNode[]} [expressions] - Array of expressions. + * @property {string} [flags] - Regex flags. + * @property {boolean} [generator] - True if function is a generator. + * @property {ASTNode} [id] - Identifier node. + * @property {ASTNode} [imported] - Imported node. + * @property {ASTNode} [init] - Initializer node. + * @property {boolean} [isEmpty] - True when the node is set for deletion but should be replaced with an Empty Statement instead. + * @property {boolean} [isMarked] - True when the node has already been marked for replacement or deletion. + * @property {ASTNode} [key] - Key node in properties. + * @property {string} [kind] - Kind of declaration (e.g., 'const'). + * @property {ASTNode} [label] - Label node. + * @property {Object[]} [leadingComments] - Leading comments. + * @property {number[]} [lineage] - Only present if detailed: true. Array of scopeIds representing the ancestry of this node's scope. + * @property {ASTNode} [left] - Left side of assignment or binary expression. + * @property {ASTNode} [local] - Local node. + * @property {boolean} [method] - True if method. + * @property {string} [name] - Name of identifier. + * @property {string} [operator] - Operator string. + * @property {ASTNode} [object] - Object node. + * @property {string} [pattern] - Pattern string. + * @property {ASTNode[]} [params] - Function parameters. + * @property {boolean} [prefix] - True if prefix operator. + * @property {ASTNode} [property] - Property node. + * @property {ASTNode[]} [properties] - Array of property nodes. + * @property {ASTNode[]} [quasis] - Template literal quasis. + * @property {number[]} [range] - [start, end] positions in source. + * @property {string} [raw] - Raw source string. + * @property {ASTNode} [regex] - Regex node. + * @property {ASTNode[]} [references] - Only present on identifier and declaration nodes (detailed: true). + * @property {ASTNode} [right] - Right side of assignment or binary expression. + * @property {ASTScope} [scope] - Only present if detailed: true. The lexical scope for this node. + * @property {number} [scopeId] - For nodes which are also a scope's block. + * @property {string} [scriptHash] - Used for caching/iteration in some utilities. + * @property {boolean} [shorthand] - True if shorthand property. + * @property {ASTNode} [source] - Source node. + * @property {string} [sourceType] - Source type (e.g., 'module'). + * @property {ASTNode[]} [specifiers] - Import/export specifiers. + * @property {boolean} [static] - True if static property. + * @property {number} [start] - Start position in source. + * @property {ASTNode} [superClass] - Superclass node. + * @property {boolean} [tail] - True if tail in template literal. + * @property {ASTNode} [test] - Test node (e.g., in IfStatement). + * @property {Object[]} [tokens] - Tokens array. + * @property {Object[]} [trailingComments] - Trailing comments. + * @property {ASTNode} [update] - Update node. + * @property {ASTNode|string|number|boolean} [value] - Value of the node. + * @property {string|function} [src] - The source code for the node, or a getter function if includeSrc is true. */ class ASTNode {} /** * @typedef ASTScope * @extends Scope - * @property {ASTNode} block - * @property {ASTScope[]} childScopes - * @property {number} scopeId - * @property {string} type + * @property {ASTNode} block - The AST node that is the block for this scope. + * @property {ASTScope[]} childScopes - Array of child scopes. + * @property {number} scopeId - Unique id for this scope. + * @property {string} type - Scope type. */ class ASTScope extends Scope {} diff --git a/tests/arborist.test.js b/tests/arborist.test.js index ba5ca26..849878f 100644 --- a/tests/arborist.test.js +++ b/tests/arborist.test.js @@ -156,4 +156,87 @@ const c = 3;`; arb.applyChanges(); assert.equal(arb.script, expected); }); +}); + +describe('Arborist edge case tests', () => { + it('Preserves comments when replacing a non-root node', () => { + const code = `const a = 1; // trailing\nconst b = 2;`; + const expected = `const a = 1;\n// trailing\nconst b = 3;`; + const arb = new Arborist(code); + const bDecl = arb.ast.find(n => n.type === 'VariableDeclarator' && n.id.name === 'b'); + arb.markNode(bDecl.init, {type: 'Literal', value: 3, raw: '3'}); + arb.applyChanges(); + assert.equal(arb.script, expected); + }); + + it('Deleting the only element in an array leaves parent valid', () => { + const code = `const a = [42];`; + const expected = `const a = [];`; + const arb = new Arborist(code); + const literal = arb.ast.find(n => n.type === 'Literal'); + arb.markNode(literal); + arb.applyChanges(); + assert.equal(arb.script, expected); + }); + + it('Multiple changes in a single pass (replace and delete siblings)', () => { + const code = `let a = 1, b = 2, c = 3;`; + const expected = `let a = 10, c = 3;`; + const arb = new Arborist(code); + const bDecl = arb.ast.find(n => n.type === 'VariableDeclarator' && n.id.name === 'b'); + const aDecl = arb.ast.find(n => n.type === 'VariableDeclarator' && n.id.name === 'a'); + arb.markNode(bDecl); // delete b + arb.markNode(aDecl.init, {type: 'Literal', value: 10, raw: '10'}); // replace a's value + arb.applyChanges(); + assert.equal(arb.script, expected); + }); + + it('Deeply nested node replacement', () => { + const code = `if (a) { if (b) { c(); } }`; + const expected = `if (a) { + if (b) { + d(); + } +}`; + const arb = new Arborist(code); + const cCall = arb.ast.find(n => n.type === 'Identifier' && n.name === 'c'); + arb.markNode(cCall, {type: 'Identifier', name: 'd'}); + arb.applyChanges(); + assert.equal(arb.script, expected); + }); + + it('Multiple comments on a node being deleted', () => { + const code = `// lead1\n// lead2\nconst a = 1; // trail1\n// trail2\nconst b = 2;`; + const expected = `// lead1\n// lead2\nconst a = 1; // trail1\n // trail2`; + const arb = new Arborist(code); + const bDecl = arb.ast.find(n => n.type === 'VariableDeclaration' && n.declarations[0].id.name === 'b'); + arb.markNode(bDecl); + arb.applyChanges(); + assert.equal(arb.script.trim(), expected.trim()); + }); + + it('Marking the same node for deletion and replacement only applies one change', () => { + const code = `let x = 1;`; + const expected = `let x = 2;`; + const arb = new Arborist(code); + const literal = arb.ast.find(n => n.type === 'Literal'); + arb.markNode(literal, {type: 'Literal', value: 2, raw: '2'}); + arb.markNode(literal); // Should not delete after replacement + arb.applyChanges(); + assert.equal(arb.script, expected); + }); + + it('AST is still valid and mutable after applyChanges', () => { + const code = `let y = 5;`; + const arb = new Arborist(code); + const literal = arb.ast.find(n => n.type === 'Literal'); + arb.markNode(literal, {type: 'Literal', value: 10, raw: '10'}); + arb.applyChanges(); + assert.equal(arb.script, 'let y = 10;'); // Validate the change was applied + // Now change again + const newLiteral = arb.ast.find(n => n.type === 'Literal'); + arb.markNode(newLiteral, {type: 'Literal', value: 20, raw: '20'}); + arb.applyChanges(); + assert.equal(arb.script, 'let y = 20;'); + }); }); \ No newline at end of file diff --git a/tests/functionality.test.js b/tests/functionality.test.js index 2914f7a..71dfa49 100644 --- a/tests/functionality.test.js +++ b/tests/functionality.test.js @@ -104,12 +104,12 @@ describe('Functionality tests', () => { let unparsedError = ''; let parsedError = ''; try { - unparsedAst = generateFlatAST(code, {alernateSourceTypeOnFailure: false}); + unparsedAst = generateFlatAST(code, {alternateSourceTypeOnFailure: false}); } catch (e) { unparsedError = e.message; } try { - parsedAst = generateFlatAST(code, {alernateSourceTypeOnFailure: true}); + parsedAst = generateFlatAST(code, {alternateSourceTypeOnFailure: true}); } catch (e) { parsedError = e.message; } @@ -121,7 +121,7 @@ describe('Functionality tests', () => { let result; const expectedResult = []; try { - result = generateFlatAST(code, {alernateSourceTypeOnFailure: false}); + result = generateFlatAST(code, {alternateSourceTypeOnFailure: false}); } catch (e) { result = e.message; }